diff --git a/.claude/skills/perf-host-optimization/SKILL.md b/.claude/skills/perf-host-optimization/SKILL.md index 854303aab754..4a72dd794b83 100644 --- a/.claude/skills/perf-host-optimization/SKILL.md +++ b/.claude/skills/perf-host-optimization/SKILL.md @@ -194,10 +194,10 @@ grep -rl "_prepare_tp_inputs\|prepare_inputs" tests/ **Running tests:** ```bash # Run specific test file with stop-on-first-failure -pytest tests/unittest/_torch/executor/test_pytorch_model_engine.py -v -x --timeout=120 +pytest tests/unittest/_torch/executor/test_model_engine.py -v -x --timeout=120 # Run specific test method -pytest tests/unittest/_torch/executor/test_pytorch_model_engine.py::PyTorchModelEngineTestCase::test_position_id_preparation -v -x +pytest tests/unittest/_torch/executor/test_model_engine.py::ModelEngineTestCase::test_position_id_preparation -v -x ``` For the full UT-to-file mapping, see [references/hot-path-files.md](references/hot-path-files.md). diff --git a/.claude/skills/perf-host-optimization/references/hot-path-files.md b/.claude/skills/perf-host-optimization/references/hot-path-files.md index d50696bfe93b..3546e030d7e0 100644 --- a/.claude/skills/perf-host-optimization/references/hot-path-files.md +++ b/.claude/skills/perf-host-optimization/references/hot-path-files.md @@ -114,7 +114,8 @@ When a top-level profiled function shows >80% time in a single sub-call, drill d | Modified File | Primary Test File | Key Tests | |---|---|---| -| `model_engine.py` | `tests/unittest/_torch/executor/test_pytorch_model_engine.py` | `test_position_id_preparation`, `test_pad_generation_requests`, `test_warmup` | +| `model_engine.py` | `tests/unittest/_torch/executor/test_model_engine.py` | `test_position_id_preparation`, `test_warmup` | +| `cuda_graph_runner.py` | `tests/unittest/_torch/executor/test_cuda_graph_runner_gpu.py` | `test_pad_generation_requests` | | `sampler.py` | `tests/unittest/_torch/executor/test_sampler.py` | Sampler-related tests | | `py_executor.py` | `tests/unittest/_torch/executor/test_py_executor.py` | Executor loop tests | | `resource_manager.py` | `tests/unittest/_torch/executor/test_resource_manager.py` | KV cache tests | diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 926b6dac98f8..72e66b423179 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -530,7 +530,6 @@ common-files: &common_files | tests/unittest/_torch/executor/test_chunked_logits.py | tests/unittest/_torch/executor/test_executor_request_queue.py | tests/unittest/_torch/executor/test_overlap_scheduler.py | - tests/unittest/_torch/executor/test_pytorch_model_engine.py | tests/unittest/_torch/executor/test_resource_manager.py | tests/unittest/_torch/executor/test_router_dealer_ipc.py | tests/unittest/_torch/helpers.py | @@ -1259,7 +1258,6 @@ legacy-files: &legacy_files | tests/unittest/_torch/executor/test_chunked_logits.py | tests/unittest/_torch/executor/test_executor_request_queue.py | tests/unittest/_torch/executor/test_overlap_scheduler.py | - tests/unittest/_torch/executor/test_pytorch_model_engine.py | tests/unittest/_torch/executor/test_resource_manager.py | tests/unittest/_torch/executor/test_router_dealer_ipc.py | tests/unittest/_torch/helpers.py | diff --git a/legacy-files.txt b/legacy-files.txt index d4c3ad70203a..bdf22b62a3ff 100644 --- a/legacy-files.txt +++ b/legacy-files.txt @@ -522,7 +522,6 @@ tests/unittest/_torch/debugger/test_debugger_addon.py tests/unittest/_torch/executor/test_chunked_logits.py tests/unittest/_torch/executor/test_executor_request_queue.py tests/unittest/_torch/executor/test_overlap_scheduler.py -tests/unittest/_torch/executor/test_pytorch_model_engine.py tests/unittest/_torch/executor/test_resource_manager.py tests/unittest/_torch/executor/test_router_dealer_ipc.py tests/unittest/_torch/helpers.py diff --git a/pyproject.toml b/pyproject.toml index 64bde66c7f68..005a02ea12ca 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -587,7 +587,6 @@ exclude = [ "tests/unittest/_torch/executor/test_chunked_logits.py", "tests/unittest/_torch/executor/test_executor_request_queue.py", "tests/unittest/_torch/executor/test_overlap_scheduler.py", - "tests/unittest/_torch/executor/test_pytorch_model_engine.py", "tests/unittest/_torch/executor/test_resource_manager.py", "tests/unittest/_torch/executor/test_router_dealer_ipc.py", "tests/unittest/_torch/helpers.py", diff --git a/ruff-legacy.toml b/ruff-legacy.toml index e65637d3e565..7830a8226867 100644 --- a/ruff-legacy.toml +++ b/ruff-legacy.toml @@ -539,7 +539,6 @@ include = [ "tests/unittest/_torch/executor/test_chunked_logits.py", "tests/unittest/_torch/executor/test_executor_request_queue.py", "tests/unittest/_torch/executor/test_overlap_scheduler.py", - "tests/unittest/_torch/executor/test_pytorch_model_engine.py", "tests/unittest/_torch/executor/test_resource_manager.py", "tests/unittest/_torch/executor/test_router_dealer_ipc.py", "tests/unittest/_torch/helpers.py", diff --git a/tests/integration/defs/.test_durations b/tests/integration/defs/.test_durations index 279a32ca1cb2..82c292c30670 100644 --- a/tests/integration/defs/.test_durations +++ b/tests/integration/defs/.test_durations @@ -1110,8 +1110,8 @@ "test_unittests.py::test_unittests_v2[unittest/_torch/executor/test_overlap_scheduler.py]": 449.18887196765496, "unittest/_torch/executor/test_profile_endpoints.py": 27.268500667556744, "test_unittests.py::test_unittests_v2[unittest/_torch/executor/test_profile_endpoints.py]": 27.268500667556744, - "unittest/_torch/executor/test_pytorch_model_engine.py -k \"test_spec_decode_graph_step_gather_kernel\"": 23.081429347826088, - "test_unittests.py::test_unittests_v2[unittest/_torch/executor/test_pytorch_model_engine.py -k \"test_spec_decode_graph_step_gather_kernel\"]": 23.081429347826088, + "unittest/_torch/executor/engine/runners/decoder/test_inputs.py -k \"test_spec_decode_graph_step_gather_kernel\"": 23.081429347826088, + "test_unittests.py::test_unittests_v2[unittest/_torch/executor/engine/runners/decoder/test_inputs.py -k \"test_spec_decode_graph_step_gather_kernel\"]": 23.081429347826088, "unittest/_torch/executor/test_resource_manager.py": 31.571532684283728, "test_unittests.py::test_unittests_v2[unittest/_torch/executor/test_resource_manager.py]": 31.571532684283728, "unittest/_torch/executor/test_seq_slot_sizing.py": 27.570640750670243, diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index deefc7a15ed3..4727ba8b8e40 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -342,7 +342,7 @@ l0_b200: # One-model speculative decoding step-copy kernels and the engine's decode # step that runs them (SM 100; they skip on other GPUs). - unittest/_torch/kernels/test_spec_step_copies.py - - unittest/_torch/executor/test_pytorch_model_engine.py -k "test_spec_decode_graph_step_gather_kernel" + - unittest/_torch/executor/engine/runners/decoder/test_inputs.py -k "test_spec_decode_graph_step_gather_kernel" # Kimi K3 DSpark drafter attention (CuTe DSL, SM 100) and its modeling_v2 catalog entries, plus the sm_100 cells of # fused_qk_norm_rope, which the drafter calls. - unittest/_torch/kernels/kimi_k3/test_k3_drafter_attn.py diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 25f1abb0d678..a4c6826309f6 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -146,6 +146,9 @@ unittest/_torch/attention/multi_gpu/test_mha_helix.py::test_mha_helix_distribute unittest/_torch/attention/multi_gpu/test_mha_helix.py::test_mha_helix_distributed[nccl-scenario: Scenario(dtype=torch.bfloat16, kv_cache_dtype=torch.bfloat16, num_layers=1, num_heads=16, num_kv_heads=4, head_dim=128, hidden_size=2048, rope_theta=10000.0, kv_cache_tokens_per_block=32, bias=False, batch=1, ctx_len=512, atol=0.1, rtol=0.05)] SKIP (https://nvbugs/6991647) unittest/_torch/attention/multi_gpu/test_mha_helix.py::test_mha_helix_distributed[nccl-scenario: Scenario(dtype=torch.bfloat16, kv_cache_dtype=torch.bfloat16, num_layers=1, num_heads=16, num_kv_heads=4, head_dim=128, hidden_size=2048, rope_theta=10000.0, kv_cache_tokens_per_block=32, bias=False, batch=1, ctx_len=64, atol=0.1, rtol=0.05)] SKIP (https://nvbugs/6991647) unittest/_torch/attention/test_attention_mla.py::test_attention_mla SKIP (https://nvbugs/6862662) +unittest/_torch/executor/engine/runners/decoder/test_inputs.py::DecoderInputPreparerTestCase::test_mrope_mixed_batched_and_fallback_positions_match_fallback SKIP (https://nvbugs/7003540) +unittest/_torch/executor/engine/runners/decoder/test_inputs.py::DecoderInputPreparerTestCase::test_mrope_multi_beam_generation_positions SKIP (https://nvbugs/7003540) +unittest/_torch/executor/engine/runners/decoder/test_inputs.py::DecoderInputPreparerTestCase::test_prepare_inputs_with_helix_parallelism SKIP (https://nvbugs/6818710) unittest/_torch/executor/kv_cache/test_kv_cache_v2_vocab_size.py::test_branch_only_enables_multimodal_keys_when_needed[dense-encoder_disable_noop] SKIP (https://nvbugs/6818710) unittest/_torch/executor/kv_cache/test_kv_cache_v2_vocab_size.py::test_branch_only_enables_multimodal_keys_when_needed[dense-multimodal] SKIP (https://nvbugs/6818710) unittest/_torch/executor/kv_cache/test_kv_cache_v2_vocab_size.py::test_branch_only_enables_multimodal_keys_when_needed[dense-text_only] SKIP (https://nvbugs/6818710) @@ -157,9 +160,6 @@ unittest/_torch/executor/kv_cache/test_kv_cache_v2_vocab_size.py::test_branch_on unittest/_torch/executor/kv_cache/test_kv_cache_v2_vocab_size.py::test_branch_only_enables_multimodal_keys_when_needed[qwen3_5_hybrid-text_only] SKIP (https://nvbugs/6818710) unittest/_torch/executor/kv_cache/test_kv_cache_v2_vocab_size.py::test_unresolvable_vocab_size_only_warns_when_multimodal_keys_are_needed[False] SKIP (https://nvbugs/6818710) unittest/_torch/executor/kv_cache/test_kv_cache_v2_vocab_size.py::test_unresolvable_vocab_size_only_warns_when_multimodal_keys_are_needed[True] SKIP (https://nvbugs/6818710) -unittest/_torch/executor/test_pytorch_model_engine.py::DecoderForwardTestCase::test_mrope_mixed_batched_and_fallback_positions_match_fallback SKIP (https://nvbugs/7003540) -unittest/_torch/executor/test_pytorch_model_engine.py::DecoderForwardTestCase::test_mrope_multi_beam_generation_positions SKIP (https://nvbugs/7003540) -unittest/_torch/executor/test_pytorch_model_engine.py::PyTorchModelEngineTestCase::test_prepare_inputs_with_helix_parallelism SKIP (https://nvbugs/6818710) unittest/_torch/kernels/parallel/test_fp4_linear.py::test_fp4_gemm_bias_per_backend[mnk2-cublaslt] SKIP (https://nvbugs/6581067) unittest/_torch/kernels/parallel/test_fp4_linear.py::test_fp4_gemm_bias_per_backend[mnk2-cutlass] SKIP (https://nvbugs/6581067) unittest/_torch/kernels/parallel/test_fp4_linear.py::test_fp4_gemm_bias_per_backend[mnk3-cublaslt] SKIP (https://nvbugs/6581067) diff --git a/tests/unittest/_torch/executor/engine/runners/__init__.py b/tests/unittest/_torch/executor/engine/runners/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tests/unittest/_torch/executor/engine/runners/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tests/unittest/_torch/executor/engine/runners/decoder/__init__.py b/tests/unittest/_torch/executor/engine/runners/decoder/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tests/unittest/_torch/executor/engine/runners/decoder/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tests/unittest/_torch/executor/engine/runners/decoder/test_dummy_requests.py b/tests/unittest/_torch/executor/engine/runners/decoder/test_dummy_requests.py new file mode 100644 index 000000000000..6c0049d8376c --- /dev/null +++ b/tests/unittest/_torch/executor/engine/runners/decoder/test_dummy_requests.py @@ -0,0 +1,62 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + + +import pytest +import torch +from _torch.executor.engine.runners.decoder.utils import decoder_config +from _torch.executor.engine_test_utils import DummyModel, kv_cache_manager + +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.dummy_requests import DummyRequestFactory +from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManager, ResourceManagerType +from tensorrt_llm.llmapi.llm_args import ( + DecodingBaseConfig, + DraftTargetDecodingConfig, + KvCacheConfig, + PARDDecodingConfig, +) + + +@pytest.mark.parametrize("config_cls", [DraftTargetDecodingConfig, PARDDecodingConfig]) +@pytest.mark.parametrize( + "draft_len", [None, 0, 3, 1], ids=["general", "cuda_graph_0", "cuda_graph_3", "cuda_graph_1"] +) +def test_warmup_builders_resynchronize_stale_draft_buffers( + config_cls: type[DecodingBaseConfig], draft_len: int | None +) -> None: + spec = config_cls(max_draft_len=3, speculative_model="dummy", draft_len_schedule={1: 3, 4: 2}) + config = decoder_config(spec_config=spec) + model = DummyModel(torch.half) + requests = DummyRequestFactory( + config=config, + model=model, + get_runtime_tokens_per_gen_step=spec.get_runtime_tokens_per_gen_step, + release_request_state=None, + ) + with kv_cache_manager( + model=model, + config=KvCacheConfig(max_tokens=258, use_kv_cache_manager_v2=True), + max_num_tokens=258, + ) as manager: + resource_manager = ResourceManager({ResourceManagerType.KV_CACHE_MANAGER: manager}) + batch_size = 2 + if draft_len is None: + expected_draft_len = config.max_draft_len + warmup_request = requests.create_warmup_request( + resource_manager, + num_tokens=batch_size * spec.tokens_per_gen_step, + num_gen_requests=batch_size, + ) + else: + expected_draft_len = draft_len + warmup_request = requests.create_cuda_graph_warmup_request( + resource_manager, batch_size, draft_len, max_seq_len=16 + ) + with requests.release_batch_context(warmup_request, resource_manager) as batch: + assert batch is not None + assert len(batch.generation_requests) == batch_size + expected_buffer_width = spec.get_runtime_tokens_per_gen_step(expected_draft_len) - 1 + assert all( + len(request.py_draft_tokens) == expected_buffer_width + for request in batch.generation_requests + ) diff --git a/tests/unittest/_torch/executor/engine/runners/decoder/test_graphs.py b/tests/unittest/_torch/executor/engine/runners/decoder/test_graphs.py new file mode 100644 index 000000000000..4d42ab8f70e9 --- /dev/null +++ b/tests/unittest/_torch/executor/engine/runners/decoder/test_graphs.py @@ -0,0 +1,539 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + + +import contextlib +import unittest +from collections.abc import Iterator +from types import SimpleNamespace +from unittest.mock import Mock, call, patch + +import pytest +import torch +from _torch.executor.engine.runners.decoder.utils import decoder_config, decoder_graphs +from _torch.executor.engine_test_utils import _create_request_with_tokens, _make_request_stub + +from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import CUDAGraphRunner +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.graphs import ( + DecoderCudaGraphs, + _make_single_token_context_graph_batch, +) +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.runner import DecoderRunner +from tensorrt_llm._torch.pyexecutor.engine.runners.encoder_decoder import EncoderDecoderRunner +from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManager +from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests +from tensorrt_llm.llmapi.llm_args import PrefillCudaGraphBackend +from tensorrt_llm.mapping import Mapping + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + ("runner_type", "graphs_path"), + [ + ( + DecoderRunner, + "tensorrt_llm._torch.pyexecutor.engine.runners.decoder.runner.DecoderCudaGraphs", + ), + ( + EncoderDecoderRunner, + "tensorrt_llm._torch.pyexecutor.engine.runners.encoder_decoder.EncoderDecoderCudaGraphs", + ), + ], +) +def test_cuda_graph_factory_preserves_input_preparer( + runner_type: type[DecoderRunner], graphs_path: str +) -> None: + runner = Mock() + runner._inputs = object() + runner._cuda_graphs_kwargs.side_effect = lambda: DecoderRunner._cuda_graphs_kwargs(runner) + + with patch(graphs_path) as graphs_type: + graphs = runner_type._make_cuda_graphs(runner) + + graphs_type.assert_called_once() + assert graphs_type.call_args.kwargs["inputs"] is runner._inputs + assert graphs_type.call_args.kwargs["warmup_inputs"] is runner._buffers.input_ids_cuda + assert graphs is graphs_type.return_value + + +@pytest.mark.cpu_only +class SingleTokenContextGraphBatchTestCase(unittest.TestCase): + def test_generation_only_is_identity(self) -> None: + generation = _make_request_stub(1) + batch = ScheduledRequests() + batch.generation_requests = [generation] + + graph_batch, promoted_ids = _make_single_token_context_graph_batch(batch) + + self.assertIs(graph_batch, batch) + self.assertEqual(promoted_ids, frozenset()) + + def test_eligible_batch_has_independent_lists_and_stable_order(self) -> None: + context_0 = _make_request_stub(10, prompt_len=1) + context_1 = _make_request_stub(11, prompt_len=8) + generation = _make_request_stub(12, prompt_len=16) + paused = object() + batch = ScheduledRequests() + batch.context_requests_last_chunk = [context_0, context_1] + batch.generation_requests = [generation] + batch.paused_requests = [paused] + semantic_lists = ( + batch.encoder_requests, + batch.context_requests_chunking, + batch.context_requests_last_chunk, + batch.generation_requests, + batch.paused_requests, + ) + semantic_snapshot = vars(context_1).copy() + + graph_batch, promoted_ids = _make_single_token_context_graph_batch(batch) + + self.assertIsNot(graph_batch, batch) + self.assertEqual(graph_batch.context_requests, []) + self.assertEqual(graph_batch.generation_requests, [context_0, context_1, generation]) + self.assertEqual(graph_batch.paused_requests, [paused]) + graph_lists = ( + graph_batch.encoder_requests, + graph_batch.context_requests_chunking, + graph_batch.context_requests_last_chunk, + graph_batch.generation_requests, + graph_batch.paused_requests, + ) + for semantic_list, graph_list in zip(semantic_lists, graph_lists): + self.assertIsNot(semantic_list, graph_list) + self.assertEqual(promoted_ids, frozenset({10, 11})) + self.assertEqual(vars(context_1), semantic_snapshot) + + graph_batch.generation_requests.append(object()) + self.assertEqual(batch.context_requests_last_chunk, [context_0, context_1]) + self.assertEqual(batch.generation_requests, [generation]) + + def test_structural_fallbacks_return_semantic_batch(self) -> None: + context = _make_request_stub(1) + + encoder_batch = ScheduledRequests() + encoder_batch.encoder_requests = [object()] + encoder_batch.context_requests_last_chunk = [context] + graph_batch, promoted_ids = _make_single_token_context_graph_batch(encoder_batch) + self.assertIs(graph_batch, encoder_batch) + self.assertFalse(promoted_ids) + + chunking_batch = ScheduledRequests() + chunking_batch.context_requests_chunking = [context] + graph_batch, promoted_ids = _make_single_token_context_graph_batch(chunking_batch) + self.assertIs(graph_batch, chunking_batch) + self.assertFalse(promoted_ids) + + def test_context_shape_and_mode_fallback_matrix(self) -> None: + cases = ( + ("multi_token_chunk", "context_chunk_size", 2), + ("more_context_remaining", "context_remaining_length", 2), + ("cursor_prompt_mismatch", "py_prompt_len", 5), + ("beam", "py_beam_width", 2), + ("draft", "py_draft_tokens", [9]), + ("context_only", "is_context_only_request", True), + ("disaggregated", "py_disaggregated_params", object()), + ("multimodal", "py_multimodal_data", {}), + ("multimodal_event", "py_mm_encoder_event", object()), + ) + for name, attribute, value in cases: + with self.subTest(name=name): + context = _make_request_stub(1) + setattr(context, attribute, value) + batch = ScheduledRequests() + batch.context_requests_last_chunk = [context] + + graph_batch, promoted_ids = _make_single_token_context_graph_batch(batch) + + self.assertIs(graph_batch, batch) + self.assertFalse(promoted_ids) + + def test_context_logits_use_final_token_graph_candidate(self) -> None: + context = _make_request_stub(1) + context.py_return_context_logits = True + batch = ScheduledRequests() + batch.context_requests_last_chunk = [context] + + graph_batch, promoted_ids = _make_single_token_context_graph_batch(batch) + + self.assertIsNot(graph_batch, batch) + self.assertEqual(graph_batch.generation_requests, [context]) + self.assertEqual(promoted_ids, frozenset({context.py_request_id})) + + def test_generation_only_request_in_context_list_falls_back(self) -> None: + context = _make_request_stub(1) + context.is_generation_only_request = True + batch = ScheduledRequests() + batch.context_requests_last_chunk = [context] + graph_batch, promoted_ids = _make_single_token_context_graph_batch(batch) + self.assertIs(graph_batch, batch) + self.assertFalse(promoted_ids) + + def test_generation_shape_fallback_matrix(self) -> None: + cases = ( + ("beam", "py_beam_width", 2), + ("draft", "py_draft_tokens", [9]), + ("disaggregated", "py_disaggregated_params", object()), + ) + for name, attribute, value in cases: + with self.subTest(name=name): + context = _make_request_stub(1) + generation = _make_request_stub(2) + setattr(generation, attribute, value) + batch = ScheduledRequests() + batch.context_requests_last_chunk = [context] + batch.generation_requests = [generation] + + graph_batch, promoted_ids = _make_single_token_context_graph_batch(batch) + + self.assertIs(graph_batch, batch) + self.assertFalse(promoted_ids) + + def test_mixed_one_and_two_token_contexts_fall_back_together(self) -> None: + one_token = _make_request_stub(1) + two_tokens = _make_request_stub(2) + two_tokens.context_current_position -= 1 + two_tokens.context_remaining_length = 2 + two_tokens.context_chunk_size = 2 + batch = ScheduledRequests() + batch.context_requests_last_chunk = [one_token, two_tokens] + + graph_batch, promoted_ids = _make_single_token_context_graph_batch(batch) + + self.assertIs(graph_batch, batch) + self.assertFalse(promoted_ids) + self.assertEqual(batch.context_requests_last_chunk, [one_token, two_tokens]) + + def test_mrope_delta_is_supported_by_decode_provider(self) -> None: + context = _make_request_stub(1) + context.py_mrope_position_delta = object() + batch = ScheduledRequests() + batch.context_requests_last_chunk = [context] + + graph_batch, promoted_ids = _make_single_token_context_graph_batch(batch) + + self.assertIsNot(graph_batch, batch) + self.assertEqual(graph_batch.generation_requests, [context]) + self.assertEqual(promoted_ids, frozenset({context.py_request_id})) + + def test_multimodal_context_requires_compatible_decode_token(self) -> None: + context = _make_request_stub(1) + context.py_multimodal_data = {} + batch = ScheduledRequests() + batch.context_requests_last_chunk = [context] + + graph_batch, promoted_ids = _make_single_token_context_graph_batch(batch) + self.assertIs(graph_batch, batch) + self.assertFalse(promoted_ids) + + incompatible = Mock(return_value=False) + graph_batch, promoted_ids = _make_single_token_context_graph_batch(batch, incompatible) + self.assertIs(graph_batch, batch) + self.assertFalse(promoted_ids) + incompatible.assert_called_once_with(context) + + compatible = Mock(return_value=True) + graph_batch, promoted_ids = _make_single_token_context_graph_batch(batch, compatible) + self.assertIsNot(graph_batch, batch) + self.assertEqual(graph_batch.generation_requests, [context]) + self.assertEqual(promoted_ids, frozenset({context.py_request_id})) + compatible.assert_called_once_with(context) + + def test_multimodal_pending_event_is_rechecked(self) -> None: + context = _make_request_stub(1) + context.py_multimodal_data = {} + context.py_mm_encoder_event = object() + batch = ScheduledRequests() + batch.context_requests_last_chunk = [context] + compatible = Mock(return_value=True) + + graph_batch, promoted_ids = _make_single_token_context_graph_batch(batch, compatible) + self.assertIs(graph_batch, batch) + self.assertFalse(promoted_ids) + compatible.assert_not_called() + + context.py_mm_encoder_event = None + graph_batch, promoted_ids = _make_single_token_context_graph_batch(batch, compatible) + self.assertIsNot(graph_batch, batch) + self.assertEqual(promoted_ids, frozenset({context.py_request_id})) + compatible.assert_called_once_with(context) + + def test_multimodal_decode_compatibility_uses_final_prompt_token(self) -> None: + graphs = decoder_graphs(decoder_config()) + graphs.model = SimpleNamespace( + config=SimpleNamespace(vocab_size=100), + mm_token_ids=torch.tensor([99], dtype=torch.int32), + ) + request = _create_request_with_tokens([11, 99, 22], 1) + + request.context_current_position = 2 + self.assertTrue(graphs._is_final_multimodal_context_decode_compatible(request)) + + request.context_current_position = 1 + self.assertFalse(graphs._is_final_multimodal_context_decode_compatible(request)) + + graphs.model.mm_token_ids = None + request = _create_request_with_tokens([11, 100], 2) + request.context_current_position = 1 + self.assertFalse(graphs._is_final_multimodal_context_decode_compatible(request)) + + graphs.model.mm_token_ids = torch.tensor([99], dtype=torch.int32) + graphs.model.model_config = SimpleNamespace( + pretrained_config=SimpleNamespace(rope_scaling={"type": "mrope"}) + ) + graphs = decoder_graphs(decoder_config(use_mrope=True), model=graphs.model) + request = _create_request_with_tokens([11, 22], 3) + request.context_current_position = 1 + request.py_multimodal_data = {"mrope_config": {}} + self.assertTrue(graphs._is_final_multimodal_context_decode_compatible(request)) + + request.py_multimodal_data["multimodal_embedding"] = object() + self.assertFalse(graphs._is_final_multimodal_context_decode_compatible(request)) + + request.py_multimodal_data["mrope_config"]["mrope_position_deltas"] = object() + self.assertTrue(graphs._is_final_multimodal_context_decode_compatible(request)) + + def test_lora_graph_variant_selection(self) -> None: + graphs = decoder_graphs(decoder_config()) + graphs._lora = SimpleNamespace(cuda_graph_manager=object()) + graphs._force_lora_graph_for_capture = None + graphs = decoder_graphs(decoder_config(cuda_graph_specialize_lora=True), lora=graphs._lora) + request = _make_request_stub(7) + batch = ScheduledRequests() + batch.generation_requests = [request] + + self.assertFalse(graphs.use_lora_cuda_graph(batch)) + + request.lora_task_id = 42 + self.assertTrue(graphs.use_lora_cuda_graph(batch)) + + request.lora_task_id = None + graphs = decoder_graphs(decoder_config(cuda_graph_specialize_lora=False), lora=graphs._lora) + self.assertTrue(graphs.use_lora_cuda_graph(batch)) + + graphs._force_lora_graph_for_capture = False + self.assertFalse(graphs.use_lora_cuda_graph(batch)) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("fails", [False, True], ids=["success", "failure"]) +def test_generation_capture_passes_local_speculation_state(fails: bool) -> None: + runner = decoder_graphs( + decoder_config(), cuda_graph_runner=Mock(enabled=True, is_warmup_only=False) + ) + resource_manager = Mock() + assert not runner._config.is_spec_decode + batches = {} + + def create_request( + resources: ResourceManager, batch_size: int, draft_len: int, *args: object, **kwargs: object + ) -> ScheduledRequests: + return batches.setdefault(draft_len, ScheduledRequests()) + + forward_warmup = Mock(side_effect=RuntimeError("capture failure") if fails else None) + with ( + patch.object(runner, "_get_graphs_to_capture", return_value=[(1, 0), (1, 2)]), + patch.object( + runner._dummy_requests, "create_cuda_graph_warmup_request", side_effect=create_request + ), + patch.object( + runner._dummy_requests, + "release_batch_context", + side_effect=lambda request, resources: contextlib.nullcontext(request), + ), + patch.object(runner, "_forward_warmup", forward_warmup), + patch.object(torch.cuda, "synchronize"), + ): + if fails: + with pytest.raises(RuntimeError, match="capture failure"): + runner._capture_generation_cuda_graphs(resource_manager) + else: + runner._capture_generation_cuda_graphs(resource_manager) + + expected_calls = [ + call(batches[2], resource_manager, enable_spec_decode=True, runtime_draft_len=2) + ] + if not fails: + expected_calls.append( + call(batches[0], resource_manager, enable_spec_decode=False, runtime_draft_len=0) + ) + assert forward_warmup.call_args_list == expected_calls + assert runner.force_lora_graph_for_capture is None + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("schedule", [{4: 4, 8: 2, 16: 1}, {4: 2}, {16: 4}, {4: 0}]) +def test_dynamic_warmup_graphs_match_replay_schedule(schedule: dict[int, int]) -> None: + from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.graphs import ( + compute_dynamic_draft_len_mapping, + ) + from tensorrt_llm._torch.speculative.utils import get_draft_len_for_batch_size + + batch_sizes = [1, 2, 4, 8, 16] + mode = SimpleNamespace(support_dynamic_draft_len=lambda: True) + runner = decoder_graphs( + decoder_config( + is_spec_decode=False, + spec_config=SimpleNamespace(draft_len_schedule=schedule, spec_dec_mode=mode), + cuda_graph_batch_sizes=batch_sizes, + original_max_draft_len=max(schedule.values()), + ) + ) + runner._dynamic_draft_len_mapping = compute_dynamic_draft_len_mapping(runner._config) + graphs = runner._get_graphs_to_capture(batch_sizes) + replay = object.__new__(CUDAGraphRunner) + replay.clear = Mock() + replay.supported_batch_sizes = batch_sizes + replay.dynamic_draft_len_mapping = runner._dynamic_draft_len_mapping + + reachable_graphs = set() + for batch_size in range(1, max(batch_sizes) + 1): + draft_len = get_draft_len_for_batch_size(schedule, batch_size, max(schedule.values())) + padded_size = replay._round_up_batch_size_with_draft_len(batch_size, draft_len) + assert padded_size >= batch_size + reachable_graphs.add((padded_size, draft_len)) + + assert set(graphs) == reachable_graphs + assert len(graphs) == len(reachable_graphs) + + +class TestWarmupCleanup(unittest.TestCase): + @pytest.mark.cpu_only + def test_cuda_graph_metrics_exclude_piecewise_stages(self) -> None: + """Verify generation graph timing excludes prefill and includes LoRA warmup cleanup.""" + graphs = DecoderCudaGraphs( + cuda_graph_runner=SimpleNamespace(enabled=True, is_warmup_only=False), + breakable_cuda_graph_runner=None, + config=SimpleNamespace( + enable_autotuner=True, prefill_cuda_graph_backend=PrefillCudaGraphBackend.PIECEWISE + ), + mapping=Mapping(), + model=SimpleNamespace(modules=lambda: []), + sparse_attention_config=None, + metrics={}, + dummy_requests=None, + agreement=None, + metadata=None, + get_runtime_tokens_per_gen_step=None, + dynamic_draft_len_mapping=None, + lora=SimpleNamespace(cuda_graph_manager=object()), + forward_warmup=None, + torch_compile_bypass=contextlib.nullcontext, + context_promotion_supported=True, + ) + resource_manager = ResourceManager({}) + events = [] + + @contextlib.contextmanager + def record_metric_scope(metric_name: str, _metrics: dict[str, float]) -> Iterator[None]: + events.append(("enter", metric_name)) + try: + yield + finally: + events.append(("exit", metric_name)) + + autotuner = SimpleNamespace( + cache_pp_recv=lambda: events.append(("lora", "cache_pp_recv")), + cache_pp_send=lambda: events.append(("lora", "cache_pp_send")), + clean_pp_flag=lambda: events.append(("lora", "clean_pp_flag")), + ) + lora_cleanup_events = [ + ("lora", "cache_pp_recv"), + ("lora", "cache_pp_send"), + ("lora", "clean_pp_flag"), + ("exit", "lora_autotune"), + ("exit", "gen_cuda_graph_warmup_seconds"), + ] + + with ( + patch( + "tensorrt_llm._torch.pyexecutor.engine.runners.decoder.graphs.autotune", + side_effect=lambda **_: record_metric_scope("lora_autotune", {}), + ), + patch( + "tensorrt_llm._torch.pyexecutor.engine.runners.decoder.graphs.AutoTuner.get", + return_value=autotuner, + ), + patch( + "tensorrt_llm._torch.modules.linear.flashinfer_mxfp8_decode_graph_capture", + side_effect=contextlib.nullcontext, + ), + patch( + "tensorrt_llm._torch.pyexecutor.engine.runners.decoder.graphs.timing_metric", + side_effect=record_metric_scope, + ), + patch.object( + graphs, + "_capture_generation_cuda_graphs", + side_effect=lambda _: events.append(("stage", "generation")), + ) as generation, + patch.object( + graphs, + "_capture_additional_cuda_graphs", + side_effect=lambda _: events.append(("stage", "additional")), + ), + patch.object( + graphs, + "_capture_prefill_cuda_graphs", + side_effect=lambda _: events.append(("stage", "piecewise")), + ) as piecewise, + ): + graphs.run_cuda_graph_warmup(resource_manager) + + self.assertEqual( + events, + [ + ("enter", "gen_cuda_graph_capture_seconds"), + ("stage", "generation"), + ("stage", "additional"), + ("exit", "gen_cuda_graph_capture_seconds"), + ("stage", "piecewise"), + ], + ) + + events.clear() + piecewise.reset_mock() + graphs.cuda_graph_runner.is_warmup_only = True + graphs.run_cuda_graph_warmup(resource_manager) + + self.assertEqual( + events, + [ + ("enter", "gen_cuda_graph_warmup_seconds"), + ("enter", "lora_autotune"), + ("stage", "generation"), + ("stage", "additional"), + *lora_cleanup_events, + ], + ) + piecewise.assert_not_called() + + events.clear() + generation.side_effect = RuntimeError("warmup failed") + with self.assertRaisesRegex(RuntimeError, "warmup failed"): + graphs.run_cuda_graph_warmup(resource_manager) + self.assertEqual( + events, + [ + ("enter", "gen_cuda_graph_warmup_seconds"), + ("enter", "lora_autotune"), + *lora_cleanup_events, + ], + ) + piecewise.assert_not_called() + + events.clear() + generation.reset_mock() + graphs.cuda_graph_runner.enabled = False + graphs._config.prefill_cuda_graph_backend = PrefillCudaGraphBackend.DISABLED + graphs.run_cuda_graph_warmup(resource_manager) + self.assertEqual( + events, + [ + ("enter", "gen_cuda_graph_warmup_seconds"), + ("enter", "lora_autotune"), + *lora_cleanup_events, + ], + ) + generation.assert_not_called() + piecewise.assert_not_called() diff --git a/tests/unittest/_torch/executor/engine/runners/decoder/test_inputs.py b/tests/unittest/_torch/executor/engine/runners/decoder/test_inputs.py new file mode 100644 index 000000000000..b2d006378032 --- /dev/null +++ b/tests/unittest/_torch/executor/engine/runners/decoder/test_inputs.py @@ -0,0 +1,1300 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + + +from __future__ import annotations + +import contextlib +import unittest +from types import SimpleNamespace +from typing import TYPE_CHECKING +from unittest.mock import Mock, patch + +import pytest +import torch +from _torch.executor.engine.runners.decoder.utils import decoder_config, input_preparer +from _torch.executor.engine_test_utils import ( + DummyModel, + _create_request, + _create_request_with_tokens, + _make_request_stub, +) +from _torch.executor.engine_test_utils import kv_cache_manager as managed_kv_cache +from utils.util import skip_pre_blackwell_unittest, skip_ray + +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionMetadata, + AttentionRuntimeFeatures, +) +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata +from tensorrt_llm._torch.attention.backends.vanilla import VanillaAttention +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.inputs import ( + DecoderInputBuffers, + DecoderInputPreparer, + HostInputs, + _get_context_prompt_lookahead_token, + _zero_fill_row_spans, +) +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.metadata import DecoderMetadata +from tensorrt_llm._torch.pyexecutor.engine.runners.interface import PreparedInputs +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest +from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManager, ResourceManagerType +from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests +from tensorrt_llm._torch.speculative.interface import INVALID_PROMPT_LOOKAHEAD_TOKEN +from tensorrt_llm._torch.speculative.spec_sampler_base import SampleStateTensorsSpec +from tensorrt_llm.bindings import SamplingConfig +from tensorrt_llm.inputs.multimodal import ( + MULTIMODAL_ENCODER_ITEM_METADATA_KEY, + MultimodalItemMetadata, +) +from tensorrt_llm.llmapi import ( + MTPDecodingConfig, + NGramDecodingConfig, + SADecodingConfig, + SamplingParams, +) +from tensorrt_llm.llmapi.llm_args import KvCacheConfig +from tensorrt_llm.mapping import CpType, Mapping + +if TYPE_CHECKING: + from collections.abc import Iterator + from typing import Any + + from tensorrt_llm._torch.kernels.spec_step_copies.op import StepInputGather + from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.config import DecoderRunnerConfig + from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 + from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig + + +def _input_preparer(*, spec_config: DecodingBaseConfig | None = None) -> DecoderInputPreparer: + return input_preparer(decoder_config(is_spec_decode=False, spec_config=spec_config), model=None) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + ("enabled", "fused"), [(None, False), (False, True), (True, False), (True, True)] +) +def test_input_preparer_configured_confidence_flags(enabled: bool | None, fused: bool) -> None: + spec_config = ( + None + if enabled is None + else SimpleNamespace( + enable_confidence_scheduling=enabled, enable_fused_confidence_scheduler=fused + ) + ) + config = SimpleNamespace(is_spec_decode=False, spec_config=spec_config) + inputs = DecoderInputPreparer( + buffers=None, + config=config, + mapping=Mapping(), + dist=None, + model=None, + get_runtime_tokens_per_gen_step=None, + iter_states={}, + lora=None, + use_beam_search=False, + steady_gen_cache_supported=True, + ) + + assert inputs._config is config + assert inputs._dspark_confidence_enabled is bool(enabled) + assert inputs._dspark_trims_submitted_tokens is bool(enabled) + assert inputs._dspark_device_windows is bool(enabled and fused) + assert inputs.previous_verify_lens == ([] if enabled else None) + assert inputs.has_previous_device_draft is False + + +@pytest.mark.cpu_only +def test_input_preparer_without_config_keeps_confidence_disabled() -> None: + inputs = DecoderInputPreparer( + buffers=None, + config=None, + mapping=Mapping(), + dist=None, + model=None, + get_runtime_tokens_per_gen_step=None, + iter_states={}, + lora=None, + use_beam_search=False, + steady_gen_cache_supported=True, + ) + + assert inputs._config is None + assert inputs._dspark_confidence_enabled is False + assert inputs._dspark_trims_submitted_tokens is False + assert inputs._dspark_device_windows is False + assert inputs.previous_verify_lens is None + assert inputs.has_previous_device_draft is False + + +@contextlib.contextmanager +def input_preparer_and_kv_cache( + config: DecoderRunnerConfig | None = None, + *, + spec_config: DecodingBaseConfig | None = None, + model: DummyModel | None = None, + max_tokens: int = 258, + max_batch_size: int = 13, +) -> Iterator[tuple[DecoderInputPreparer, KVCacheManagerV2]]: + config = config or decoder_config(spec_config=spec_config) + model = ( + model + if model is not None + else DummyModel(config.dtype, rope_scaling={"type": "mrope"} if config.use_mrope else None) + ) + inputs = input_preparer( + config, + model=model, + buffers=DecoderInputBuffers.for_config(config, config.max_beam_width > 1), + ) + with managed_kv_cache( + model=model, max_tokens=max_tokens, max_batch_size=max_batch_size + ) as manager: + yield inputs, manager + + +@pytest.mark.cpu_only +class TestCachedKvTokenLogging(unittest.TestCase): + def setUp(self) -> None: + self.inputs = _input_preparer() + self.real = SimpleNamespace(is_cuda_graph_dummy=False) + self.padding = SimpleNamespace(is_cuda_graph_dummy=True) + + def check_counts( + self, + counts: list[int], + groups: tuple[tuple[list[SimpleNamespace], int], ...], + expected: list[int], + padding: int, + ) -> None: + self.inputs.iter_states["cached_kv_tokens"] = sum(counts) + self.inputs.record_cached_kv_tokens_per_req(counts, groups) + self.assertEqual( + self.inputs.iter_states, + { + "cached_kv_tokens": sum(counts), + "cached_kv_tokens_per_req": expected, + "cached_kv_tokens_cuda_graph_padding": padding, + }, + ) + counts[:] = [999] + self.assertEqual(self.inputs.iter_states["cached_kv_tokens_per_req"], expected) + + def test_mtp_padding_before_generation(self) -> None: + # Input packing is context, extend (including padding), gen. + self.check_counts( + [10, 100, 7, 200], + (([self.real], 1), ([self.real, self.padding], 1), ([self.real], 1)), + [10, 100, 200], + 7, + ) + + def test_beams_preserved_after_context_and_padding(self) -> None: + self.check_counts( + [10, 100, 101, 7, 7, 200, 201], + (([self.real], 1), ([self.real, self.padding, self.real], 2)), + [10, 100, 101, 200, 201], + 14, + ) + + def test_attention_dp_dummy_is_not_graph_padding(self) -> None: + adp_dummy = SimpleNamespace( + is_dummy=True, is_attention_dp_dummy=True, is_cuda_graph_dummy=False + ) + self.check_counts([10, 1, 7], (([self.real, adp_dummy, self.padding], 1),), [10, 1], 7) + + def test_padding_only(self) -> None: + self.check_counts([7, 7], (([self.padding, self.padding], 1),), [], 14) + + def test_group_row_count_must_match_counts(self) -> None: + for counts, groups in ( + ([100, 200], (([self.real], 1),)), + ([100, 101, 200], (([self.real, self.real], 2),)), + ): + with self.subTest(counts=counts, groups=groups): + with self.assertRaises(AssertionError): + self.inputs.record_cached_kv_tokens_per_req(counts, groups) + + def test_steady_batch_clears_previous_padding(self) -> None: + self.inputs.iter_states["cached_kv_tokens_cuda_graph_padding"] = 7 + self.check_counts([100, 200], (), [100, 200], 0) + + def test_empty_batch_clears_previous_values(self) -> None: + self.inputs.iter_states["cached_kv_tokens_per_req"] = [100] + self.inputs.iter_states["cached_kv_tokens_cuda_graph_padding"] = 7 + self.check_counts([], (([], 1),), [], 0) + + +@pytest.mark.cpu_only +def test_context_prompt_lookahead_stops_at_prompt_boundary() -> None: + """An internal chunk hands back the next prompt token; the last one does not. + + The boundary is read off the request's live chunking state rather than a + Python-side prompt length, which a preemption leaves stale. + """ + request = _create_request_with_tokens([10, 11, 12, 13, 14], 1) + + request.context_chunk_size = 2 + assert _get_context_prompt_lookahead_token(request, 2) == 12 + + # A reused prefix advances the live cursor while py_prompt_len remains the + # original prompt length. Internal chunk [2, 4) still looks up token 4. + request.context_current_position = 2 + request.context_chunk_size = 2 + assert _get_context_prompt_lookahead_token(request, 4) == 14 + + # Final chunk [2, 5): nothing follows in the prompt, so the drafter falls + # back to the token the target just sampled. + request.context_chunk_size = 3 + assert _get_context_prompt_lookahead_token(request, 5) == INVALID_PROMPT_LOOKAHEAD_TOKEN + + +@pytest.mark.cpu_only +def test_zero_fill_row_spans_merges_unsorted_and_adjacent_spans() -> None: + buffer = torch.arange(1, 13, dtype=torch.int32) + _zero_fill_row_spans(buffer, [(9, 11), (2, 4), (4, 5), (7, 8)]) + assert buffer.tolist() == [1, 2, 0, 0, 0, 6, 7, 0, 9, 0, 0, 12] + + +@pytest.mark.cpu_only +def test_zero_fill_row_spans_without_spans_is_noop() -> None: + buffer = torch.arange(4, dtype=torch.int32) + _zero_fill_row_spans(buffer, []) + assert buffer.tolist() == [0, 1, 2, 3] + + +class DecoderInputPreparerTestCase(unittest.TestCase): + def test_prepare_inputs_returns_steady_gen_fast_inputs(self) -> None: + inputs = _input_preparer() + inputs.guided_decoder = Mock() + inputs._can_use_steady_gen_fast_prepare = Mock(return_value=True) + fast_inputs = PreparedInputs({"prepared": True}) + inputs._apply_steady_gen_fast_prepare = Mock(return_value=fast_inputs) + batch = ScheduledRequests() + new_tokens = torch.zeros(1) + + result = inputs.prepare_inputs( + batch, + Mock(), + Mock(), + new_tensors_device=SimpleNamespace(new_tokens=new_tokens), + enable_spec_decode=False, + runtime_draft_len=3, + is_dummy=True, + ) + + self.assertIs(result.kwargs, fast_inputs.kwargs) + self.assertEqual(result.runtime_draft_len, 3) + inputs.guided_decoder.add_batch.assert_called_once_with( + batch, new_tokens=new_tokens, runtime_draft_len=3 + ) + inputs._can_use_steady_gen_fast_prepare.assert_called_once_with( + batch, new_tokens, None, None, True + ) + + def test_prepare_inputs_reports_tree_draft_length(self) -> None: + inputs = _input_preparer( + spec_config=SimpleNamespace(is_linear_tree=False, tokens_per_gen_step=8) + ) + inputs._can_use_steady_gen_fast_prepare = Mock(return_value=True) + inputs._apply_steady_gen_fast_prepare = Mock(return_value=PreparedInputs({})) + + for enable_spec_decode, expected in ((True, 7), (False, 2)): + with self.subTest(enable_spec_decode=enable_spec_decode): + result = inputs.prepare_inputs( + ScheduledRequests(), + Mock(), + Mock(), + enable_spec_decode=enable_spec_decode, + runtime_draft_len=2, + is_dummy=True, + ) + self.assertEqual(result.runtime_draft_len, expected) + + def test_steady_gen_fast_prepare_skips_dummy_passes(self) -> None: + inputs = _input_preparer() + request = _make_request_stub(7) + inputs._steady_gen_cache = { + "num_requests": 1, + "request_ids": [request.py_request_id], + } + batch = ScheduledRequests() + batch.generation_requests = [request] + new_tokens = torch.zeros(1) + + self.assertTrue( + inputs._can_use_steady_gen_fast_prepare(batch, new_tokens, None, None, False) + ) + self.assertFalse( + inputs._can_use_steady_gen_fast_prepare(batch, new_tokens, None, None, True) + ) + + def test_padding_request_input_id_rows_are_zeroed(self) -> None: + (preparer, kv_cache_manager) = self.enterContext(input_preparer_and_kv_cache()) + resource_manager = ResourceManager({ResourceManagerType.KV_CACHE_MANAGER: kv_cache_manager}) + attn_metadata = AttentionMetadata( + max_num_requests=4, max_num_tokens=32, kv_cache_manager=kv_cache_manager + ) + attn_metadata.is_cuda_graph = False + + generation = _create_request_with_tokens([50, 51, 52, 53, 54], 1) + generation.py_seq_slot = 0 + generation.py_batch_idx = 0 + padding = _create_request_with_tokens([60, 61, 62], 2) + padding.is_cuda_graph_dummy = True + padding.py_seq_slot = 1 + + graph_batch = ScheduledRequests() + graph_batch.generation_requests = [generation, padding] + new_tokens = torch.zeros((1, 4, 1), dtype=torch.int32, device="cuda") + new_tokens[0, 0, 0] = 777 + overlap_state = SimpleNamespace(new_tokens=new_tokens) + inputs = preparer + inputs._can_use_steady_gen_fast_prepare = Mock(return_value=False) + # The padding row is written by neither input-id copy, so poison the + # buffer with an out-of-vocabulary id to catch a stale read. + preparer._buffers.input_ids_cuda.fill_(1 << 30) + + prepared = inputs.prepare_inputs( + scheduled_requests=graph_batch, + kv_cache_manager=kv_cache_manager, + attn_metadata=attn_metadata, + new_tensors_device=overlap_state, + resource_manager=resource_manager, + enable_spec_decode=False, + runtime_draft_len=0, + is_dummy=False, + ) + + self.assertEqual(prepared.kwargs["input_ids"][:2].cpu().tolist(), [777, 0]) + + def test_promoted_context_precedes_speculative_overlap_generation(self) -> None: + spec_config = SADecodingConfig( + max_draft_len=1, + draft_len_schedule={1: 1}, + ) + (preparer, kv_cache_manager) = self.enterContext( + input_preparer_and_kv_cache(spec_config=spec_config) + ) + resource_manager = ResourceManager({ResourceManagerType.KV_CACHE_MANAGER: kv_cache_manager}) + attn_metadata = AttentionMetadata( + max_num_requests=4, max_num_tokens=32, kv_cache_manager=kv_cache_manager + ) + attn_metadata.is_cuda_graph = False + # A bare Mock auto-vivifies every attribute, so optional metadata + # fields read by prepare_inputs must be pinned off explicitly. + spec_metadata = Mock( + _force_non_greedy_for_capture=False, + context_prompt_lookahead_tokens=None, + ) + + context = _create_request_with_tokens([11, 22, 33, 44], 1) + context.context_current_position = 3 + context.context_chunk_size = 1 + context.py_seq_slot = 0 + # A promoted context row must ignore any stale overlap slot. + context.py_batch_idx = 3 + + generation = _create_request_with_tokens([50, 51, 52, 53, 54], 2) + generation.py_seq_slot = 1 + generation.py_batch_idx = 1 + generation.py_needs_onehot_draft_probs = True + + graph_batch = ScheduledRequests() + graph_batch.generation_requests = [context, generation] + new_tokens = torch.zeros((1, 4, 1), dtype=torch.int32, device="cuda") + new_tokens[0, 0, 0] = 999 + new_tokens[0, 1, 0] = 777 + overlap_state = SampleStateTensorsSpec( + new_tokens=new_tokens, + new_tokens_lens=torch.ones(4, dtype=torch.int32, device="cuda"), + next_draft_tokens=torch.zeros((4, 1), dtype=torch.int32, device="cuda"), + ) + + prepared = preparer.prepare_inputs( + scheduled_requests=graph_batch, + kv_cache_manager=kv_cache_manager, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + new_tensors_device=overlap_state, + resource_manager=resource_manager, + promoted_context_request_ids=frozenset({context.py_request_id}), + enable_spec_decode=True, + runtime_draft_len=0, + is_dummy=False, + ) + + inputs = prepared.kwargs + self.assertEqual(inputs["input_ids"][:2].cpu().tolist(), [44, 777]) + self.assertEqual(inputs["position_ids"][0, :2].cpu().tolist(), [3, 4]) + self.assertEqual( + attn_metadata.request_ids, [context.py_request_id, generation.py_request_id] + ) + self.assertEqual(attn_metadata.kv_cache_params.num_cached_tokens_per_seq, [3, 5]) + self.assertEqual(preparer._buffers.previous_batch_indices_cuda[:1].cpu().tolist(), [1]) + self.assertEqual(preparer._buffers.previous_pos_id_offsets_cuda[:2].cpu().tolist(), [0, 1]) + self.assertEqual(attn_metadata.num_contexts, 0) + self.assertEqual( + spec_metadata.request_ids, [context.py_request_id, generation.py_request_id] + ) + self.assertFalse(generation.py_needs_onehot_draft_probs) + spec_metadata.write_padding_onehot_draft_probs.assert_called_once_with( + [generation.py_seq_slot], 0 + ) + + def test_spec_decode_graph_step_gather_kernel_matches_torch_path(self) -> None: + """A CUDA graph decode step of a speculative engine whose overlap + gathers run as one StepInputGather launch writes what the torch path + writes.""" + from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata + from tensorrt_llm._torch.kernels.spec_step_copies import op as spec_step_copies + + if not spec_step_copies.is_supported(): + self.skipTest("the step-copy kernels run on SM 100") + max_draft_len = 3 + tokens_per_step = max_draft_len + 1 + (preparer, kv_cache_manager) = self.enterContext( + input_preparer_and_kv_cache(spec_config=SADecodingConfig(max_draft_len=max_draft_len)) + ) + inputs = preparer + buffers = preparer._buffers + step_gather = inputs._step_input_gather + self.assertIsNotNone(step_gather) + resource_manager = ResourceManager({ResourceManagerType.KV_CACHE_MANAGER: kv_cache_manager}) + + # Three requests that ran in the previous step, in slots out of order. + slots = [5, 2, 7] + requests = kv_cache_manager.add_dummy_requests( + [1, 2, 3], token_nums=[10, 17, 33], is_gen=True, max_num_draft_tokens=max_draft_len + ) + for request, slot in zip(requests, slots): + request.is_dummy_request = False + request.py_seq_slot = slot + batch = ScheduledRequests() + batch.generation_requests = requests + attn_metadata = DecoderMetadata( + config=preparer._config, + mapping=Mapping(), + model=preparer.model, + buffers=preparer._buffers, + attention_runtime_features=AttentionRuntimeFeatures(), + ).setup_attention(kv_cache_manager) + self.assertIs(type(attn_metadata), TrtllmAttentionMetadata) + graph_metadata = attn_metadata.create_cuda_graph_metadata( + len(requests), False, max_draft_len + ) + spec_metadata = Mock( + _force_non_greedy_for_capture=False, + context_prompt_lookahead_tokens=None, + ) + + # The sampler's slot stores as the previous step left them. + num_slots = 8 + generator = torch.Generator(device="cuda").manual_seed(0) + previous = SampleStateTensorsSpec( + new_tokens=torch.randint( + 0, + 1000, + (tokens_per_step, num_slots, 1), + generator=generator, + dtype=torch.int32, + device="cuda", + ), + new_tokens_lens=torch.randint( + 1, + tokens_per_step + 1, + (num_slots,), + generator=generator, + dtype=torch.int32, + device="cuda", + ), + next_draft_tokens=torch.randint( + 0, + 1000, + (num_slots, max_draft_len), + generator=generator, + dtype=torch.int32, + device="cuda", + ), + ) + + def step_inputs(step_input_gather: StepInputGather | None) -> list[torch.Tensor]: + """The buffers one step writes, from a filler they all start at.""" + inputs._step_input_gather = step_input_gather + for request, slot in zip(requests, slots): + request.py_batch_idx = slot + written = ( + buffers.input_ids_cuda, + buffers.position_ids_cuda, + buffers.draft_tokens_cuda, + buffers.previous_pos_id_offsets_cuda, + buffers.previous_kv_lens_offsets_cuda, + ) + for buffer in written: + buffer.fill_(-3) + inputs.prepare_inputs( + scheduled_requests=batch, + kv_cache_manager=kv_cache_manager, + attn_metadata=graph_metadata, + spec_metadata=spec_metadata, + new_tensors_device=previous, + resource_manager=resource_manager, + enable_spec_decode=True, + runtime_draft_len=max_draft_len, + is_dummy=False, + ) + torch.cuda.synchronize() + return [buffer.clone() for buffer in written] + + torch_path = step_inputs(None) + launched = [] + launch = step_gather.gather + + def recorded_launch(*args: object, **kwargs: object) -> bool: + launched.append(launch(*args, **kwargs)) + return launched[-1] + + with patch.object(step_gather, "gather", side_effect=recorded_launch): + kernel_path = step_inputs(step_gather) + self.assertEqual(launched, [True]) + for expected, actual in zip(torch_path, kernel_path): + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + # The previous rows' input ids came from the stores, by slot. + num_tokens = len(requests) * tokens_per_step + expected_input_ids = previous.new_tokens[:, slots, 0].t().reshape(-1) + self.assertEqual(kernel_path[0][:num_tokens].tolist(), expected_input_ids.tolist()) + + @skip_ray + def test_prepare_inputs_with_helix_parallelism(self) -> None: + """Test prepare_inputs with helix parallelism.""" + + # Create model engine with helix parallelism. + + # Provide mapping for model engine. + cp_size = 2 + cp_rank = 0 + cp_config = {"cp_type": CpType.HELIX, "tokens_per_block": 4} + mapping = Mapping( + world_size=cp_size, + tp_size=1, + pp_size=1, + cp_size=cp_size, + cp_config=cp_config, + rank=cp_rank, + ) + config = decoder_config(max_num_tokens=512) + preparer = input_preparer( + config, + model=DummyModel(torch.half), + mapping=mapping, + buffers=DecoderInputBuffers.for_config(config, False), + ) + + # Create scheduled requests with two generation requests. + scheduled_requests = ScheduledRequests() + scheduled_requests.context_requests_last_chunk = [] + prompt_lens = [20, 15] + gen_requests = [] + for idx in range(len(prompt_lens)): + req = _create_request(num_tokens=prompt_lens[idx], req_id=idx + 1) + req.py_prompt_len = prompt_lens[idx] + req.py_batch_idx = None + req.is_dummy_request = False + req.py_seq_slot = idx + req.sampling_config.beam_width = 1 + req.py_multimodal_data = {} + req.total_input_len_cp = prompt_lens[idx] * 2 + req.seqlen_this_rank_cp = prompt_lens[idx] + req.py_decoding_iter = 1 + gen_requests.append(req) + scheduled_requests.generation_requests = gen_requests + + # Create KV cache manager for attention metadata. + kv_cache_config = KvCacheConfig(max_tokens=512, enable_block_reuse=False) + kv_cache_manager = self.enterContext( + managed_kv_cache( + model=preparer.model, + max_tokens=512, + max_batch_size=4, + mapping=mapping, + config=kv_cache_config, + ) + ) + kv_cache_manager.add_dummy_requests([1, 2], prompt_lens, is_gen=True) + attn_metadata = TrtllmAttentionMetadata( + max_num_requests=4, + max_num_tokens=512, + kv_cache_manager=kv_cache_manager, + mapping=mapping, + ) + attn_metadata.is_cuda_graph = False + + # Initialize model engine buffers. + + result = preparer.prepare_inputs( + scheduled_requests=scheduled_requests, + kv_cache_manager=kv_cache_manager, + attn_metadata=attn_metadata, + enable_spec_decode=False, + runtime_draft_len=0, + is_dummy=False, + ).kwargs + + # Verify expected keys are present. + self.assertIsNotNone(result) + self.assertIn("input_ids", result) + self.assertIn("position_ids", result) + self.assertIn("attn_metadata", result) + + # Also, verify that position_ids are properly calculated. + position_ids = result["position_ids"] + self.assertIsInstance(position_ids, torch.Tensor) + expected_positions = [40, 30] + actual_positions = position_ids.squeeze(0).cpu().tolist()[:2] + self.assertEqual( + actual_positions, + expected_positions, + f"Position IDs should reflect CP allgather results. Expected: {expected_positions}, " + f"Got: {actual_positions}", + ) + + # Verify attention metadata is properly configured. + self.assertEqual(attn_metadata.request_ids, [1, 2]) + self.assertEqual(attn_metadata.prompt_lens, [20, 15]) + self.assertEqual(attn_metadata.num_contexts, 0) + + # Verify KV cache parameters + self.assertIsNotNone(attn_metadata.kv_cache_params) + self.assertTrue(attn_metadata.kv_cache_params.use_cache) + + # Verify sequence lengths are correct. + expected_seq_lens = [1, 1] + if hasattr(attn_metadata, "seq_lens") and attn_metadata.seq_lens is not None: + actual_seq_lens = attn_metadata.seq_lens.cpu().tolist() + self.assertEqual(actual_seq_lens, expected_seq_lens) + + def test_mrope_multi_beam_generation_positions(self) -> None: + """Every beam row gets its request's `past + delta` on all 3 axes.""" + beam_width = 2 + config = decoder_config( + max_batch_size=4, max_num_tokens=64, max_beam_width=beam_width, use_mrope=True + ) + preparer = input_preparer( + config, + model=DummyModel(torch.half, rope_scaling={"type": "mrope"}), + buffers=DecoderInputBuffers.for_config(config, config.max_beam_width > 1), + ) + # (num_tokens, delta): one delta arrives on the host and one is + # already a CUDA int32 scalar. + cases = ( + (4, torch.tensor([[10]], dtype=torch.int32)), + (7, torch.tensor([[-3]], dtype=torch.int32, device="cuda")), + ) + requests = [] + for seq_slot, (num_tokens, delta) in enumerate(cases): + request = LlmRequest( + request_id=seq_slot + 1, + max_new_tokens=1, + input_tokens=list(range(1, num_tokens + 1)), + sampling_config=SamplingConfig( + SamplingParams(n=beam_width, use_beam_search=True)._get_sampling_config() + ), + is_streaming=False, + ) + request.py_seq_slot = seq_slot + request.py_multimodal_data = {"mrope_config": {"mrope_position_deltas": delta}} + requests.append(request) + host = HostInputs( + use_mrope=True, + mrope_dummy_seq_slot=config.max_num_tokens, + extra_inputs=preparer._new_extra_inputs(), + context_prompt_lookahead=None, + ) + host.lists.generation_requests.extend(requests) + + batched_mrope_generation_updates: list[tuple[int, int, torch.Tensor]] = [] + preparer._pack_generation_requests( + host, + frozenset(), + None, + batched_mrope_generation_updates=batched_mrope_generation_updates, + ) + host.total_num_tokens = len(host.lists.position_ids) + preparer._copy_position_ids_to_device( + host, batched_mrope_generation_updates=batched_mrope_generation_updates + ) + + # Both deltas take the batched scalar path, one row per beam. + self.assertEqual([row for row, _, _ in batched_mrope_generation_updates], [0, 1, 2, 3]) + self.assertEqual(host.lists.mrope_position_ids, []) + # Without a previous tensor the past length is num_tokens - 1. + expected_row = [ + num_tokens - 1 + int(delta) for num_tokens, delta in cases for _ in range(beam_width) + ] + self.assertEqual(expected_row, [13, 13, 3, 3]) + torch.testing.assert_close( + preparer._buffers.mrope_position_ids_cuda[:, :, :4], + torch.tensor([[expected_row]] * 3, dtype=torch.int32, device="cuda"), + atol=0, + rtol=0, + ) + + def test_mrope_mixed_batched_and_fallback_positions_match_fallback(self) -> None: + """Mixed scalar batching matches the old per-request writes exactly.""" + config = decoder_config(max_batch_size=2, max_num_tokens=64, use_mrope=True) + preparer = input_preparer( + config, + model=DummyModel(torch.half, rope_scaling={"type": "mrope"}), + buffers=DecoderInputBuffers.for_config(config, config.max_beam_width > 1), + ) + sentinel = 123 + + def run( + delta_dtypes: tuple[torch.dtype, torch.dtype], + ) -> tuple[HostInputs, list[tuple[int, int, torch.Tensor]], torch.Tensor]: + requests = [] + for seq_slot, (num_tokens, delta, dtype) in enumerate( + zip((4, 7), (10, -3), delta_dtypes) + ): + request = LlmRequest( + request_id=seq_slot + 1, + max_new_tokens=1, + input_tokens=list(range(1, num_tokens + 1)), + sampling_config=SamplingConfig(SamplingParams()._get_sampling_config()), + is_streaming=False, + ) + request.py_seq_slot = seq_slot + request.py_multimodal_data = { + "mrope_config": { + "mrope_position_deltas": torch.tensor([[delta]], dtype=dtype, device="cuda") + } + } + requests.append(request) + host = HostInputs( + use_mrope=True, + mrope_dummy_seq_slot=config.max_num_tokens, + extra_inputs=preparer._new_extra_inputs(), + context_prompt_lookahead=None, + ) + host.lists.generation_requests.extend(requests) + batched_mrope_generation_updates: list[tuple[int, int, torch.Tensor]] = [] + preparer._pack_generation_requests( + host, + frozenset(), + None, + batched_mrope_generation_updates=batched_mrope_generation_updates, + ) + host.total_num_tokens = len(host.lists.position_ids) + preparer._buffers.mrope_position_ids_cuda.fill_(sentinel) + preparer._copy_position_ids_to_device( + host, batched_mrope_generation_updates=batched_mrope_generation_updates + ) + positions = preparer._buffers.mrope_position_ids_cuda.clone() + return host, batched_mrope_generation_updates, positions + + mixed_host, mixed_updates, mixed = run((torch.int32, torch.int64)) + _, fallback_updates, fallback = run((torch.int64, torch.int64)) + + self.assertEqual([row for row, _, _ in mixed_updates], [0]) + self.assertEqual( + [(start, end) for start, end, _ in mixed_host.lists.mrope_position_ids], [(1, 2)] + ) + self.assertEqual(fallback_updates, []) + torch.testing.assert_close(mixed, fallback, atol=0, rtol=0) + self.assertTrue(torch.all(mixed[:, :, 2:] == sentinel).item()) + + def _setup_mrope_inputs( + self, + max_num_tokens: int = 32, + spec_config: DecodingBaseConfig | None = None, + **config_changes: Any, + ) -> tuple[DecoderInputPreparer, KVCacheManagerV2, AttentionMetadata]: + config = decoder_config( + **{ + "max_num_tokens": max_num_tokens, + "spec_config": spec_config, + "use_mrope": True, + **config_changes, + } + ) + inputs, manager = self.enterContext( + input_preparer_and_kv_cache( + config, + max_tokens=max_num_tokens, + max_batch_size=4, + ) + ) + metadata = AttentionMetadata( + max_num_requests=4, max_num_tokens=max_num_tokens, kv_cache_manager=manager + ) + metadata.is_cuda_graph = False + return inputs, manager, metadata + + @staticmethod + def _make_mrope_gen_request( + num_tokens: int, req_id: int, seq_slot: int, delta: int | None + ) -> LlmRequest: + """Generation request carrying an MRoPE delta when `delta` is not None, + and no multimodal data at all otherwise (a text-only prompt).""" + request = _create_request(num_tokens, req_id) + request.py_prompt_len = num_tokens + request.py_batch_idx = None + request.py_seq_slot = seq_slot + request.sampling_config.beam_width = 1 + if delta is None: + request.py_multimodal_data = {} + else: + request.py_multimodal_data = { + "mrope_config": { + "mrope_position_deltas": torch.tensor([[delta]], dtype=torch.int32) + }, + } + return request + + def _make_mrope_spec_batch( + self, previous_batch: bool + ) -> tuple[ScheduledRequests, list[LlmRequest]]: + """Two multimodal verify groups around a text-only one, plus a CUDA graph + padding dummy, each carrying one draft token. + + Returns the scheduled batch and the multimodal requests.""" + # (num_tokens, req_id, seq_slot, delta); the middle request is text-only. + requests = [ + self._make_mrope_gen_request(4, 1, 0, 10), + self._make_mrope_gen_request(5, 2, 1, None), + self._make_mrope_gen_request(6, 3, 2, 20), + ] + for request in requests: + request.py_draft_tokens = [7] + if previous_batch: + # The overlap scheduler's previous step ran these requests at + # their current slots, which also seeded the delta cache. + request.py_batch_idx = request.py_seq_slot + if request.py_multimodal_data: + request.py_mrope_delta_cache_slot = request.py_seq_slot + # The padding dummy carries the zero delta that `add_dummy_requests` + # attaches on MRoPE models. + dummy = self._make_mrope_gen_request(2, 4, 3, 0) + dummy.is_cuda_graph_dummy = True + dummy.py_draft_tokens = [0] + + scheduled_requests = ScheduledRequests() + scheduled_requests.context_requests_last_chunk = [] + scheduled_requests.generation_requests = requests + [dummy] + return scheduled_requests, [requests[0], requests[2]] + + def test_prepare_inputs_spec_decode_adds_mrope_delta(self) -> None: + """One-model speculative verify groups get the request's MRoPE delta. + + Every token a generation step feeds is text after the prompt's vision + spans, so its three coordinates are `seq_index + delta`, as for an + ordinary generation request. Read slots stay dense, one per verify + group, and a generation-only request seeds the delta cache. + """ + preparer, kv_cache_manager, attn_metadata = self._setup_mrope_inputs( + spec_config=SADecodingConfig(max_draft_len=1) + ) + scheduled_requests, multimodal_requests = self._make_mrope_spec_batch(previous_batch=False) + spec_metadata = Mock( + _force_non_greedy_for_capture=False, + context_prompt_lookahead_tokens=None, + ) + + result = preparer.prepare_inputs( + scheduled_requests=scheduled_requests, + kv_cache_manager=kv_cache_manager, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + enable_spec_decode=True, + runtime_draft_len=1, + is_dummy=False, + ).kwargs + + # Each verify group is [last token, draft token] at past_seen + 0/1; + # the text-only request and the zero-delta dummy keep the scalar + # positions. + position_ids = result["position_ids"] + self.assertEqual(tuple(position_ids.shape), (3, 1, 8)) + expected = torch.tensor( + [[[13, 14, 4, 5, 25, 26, 1, 2]]] * 3, dtype=torch.int32, device="cuda" + ) + torch.testing.assert_close(position_ids, expected, atol=0, rtol=0) + # Slot 32 is the reserved zero slot (max_num_tokens * pp_size). + self.assertEqual(result["mrope_delta_read_seq_slots"].cpu().tolist(), [0, 32, 2, 32]) + self.assertEqual(result["mrope_delta_write_seq_slots"].cpu().tolist(), [0, 2]) + self.assertEqual( + [request.py_mrope_delta_cache_slot for request in multimodal_requests], [0, 2] + ) + + def test_prepare_inputs_spec_decode_overlap_keeps_mrope_delta(self) -> None: + """Under the overlap scheduler the accepted-token correction applied in + `preprocess_inputs` lands on top of the delta, on all three axes.""" + preparer, kv_cache_manager, attn_metadata = self._setup_mrope_inputs( + spec_config=SADecodingConfig(max_draft_len=1) + ) + scheduled_requests, _ = self._make_mrope_spec_batch(previous_batch=True) + spec_metadata = Mock( + _force_non_greedy_for_capture=False, + context_prompt_lookahead_tokens=None, + ) + # The previous step accepted 2, 1 and 2 tokens in slots 0, 1 and 2. + overlap_state = SampleStateTensorsSpec( + new_tokens=torch.zeros((2, 4, 1), dtype=torch.int32, device="cuda"), + new_tokens_lens=torch.tensor([2, 1, 2, 1], dtype=torch.int32, device="cuda"), + next_draft_tokens=torch.zeros((4, 1), dtype=torch.int32, device="cuda"), + ) + + inputs = preparer.prepare_inputs( + scheduled_requests=scheduled_requests, + kv_cache_manager=kv_cache_manager, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + new_tensors_device=overlap_state, + enable_spec_decode=True, + runtime_draft_len=1, + is_dummy=False, + ).kwargs + + # The host packs positions from the stale base (past_seen + delta). + stale = torch.tensor([[[13, 14, 4, 5, 25, 26, 1, 2]]] * 3, dtype=torch.int32, device="cuda") + torch.testing.assert_close(inputs["position_ids"], stale, atol=0, rtol=0) + self.assertEqual(inputs["mrope_delta_read_seq_slots"].cpu().tolist(), [0, 32, 2, 32]) + # The delta cache already holds both multimodal requests' deltas. + self.assertNotIn("mrope_delta_write_seq_slots", inputs) + + inputs = preparer.preprocess_inputs(inputs, enable_spec_decode=True, runtime_draft_len=1) + corrected = torch.tensor( + [[[15, 16, 5, 6, 27, 28, 1, 2]]] * 3, dtype=torch.int32, device="cuda" + ) + torch.testing.assert_close(inputs["position_ids"], corrected, atol=0, rtol=0) + + @staticmethod + def _make_mrope_ctx_request( + num_tokens: int, req_id: int, seq_slot: int, delta: int + ) -> LlmRequest: + """Single-chunk context request carrying MRoPE metadata. + + Its position on axis `a` of token `i` is `100 + 1000 * a + i`, so a + spliced span says which axis and token it came from.""" + request = _create_request(num_tokens, req_id) + request.py_prompt_len = num_tokens + request.py_batch_idx = None + request.py_seq_slot = seq_slot + request.context_current_position = 0 + request.context_chunk_size = num_tokens + request.sampling_config.beam_width = 1 + positions = ( + 100 + 1000 * torch.arange(3).view(3, 1, 1) + torch.arange(num_tokens).view(1, 1, -1) + ).to(torch.int32) + request.py_multimodal_data = { + "mrope_config": { + "mrope_position_ids": positions, + "mrope_position_deltas": torch.tensor([[delta]], dtype=torch.int32), + }, + } + return request + + @staticmethod + def _prepare_mrope_spec_inputs( + preparer: DecoderInputPreparer, + kv_cache_manager: KVCacheManagerV2, + attn_metadata: AttentionMetadata, + scheduled_requests: ScheduledRequests, + ) -> dict[str, Any]: + spec_metadata = Mock( + _force_non_greedy_for_capture=False, + context_prompt_lookahead_tokens=None, + ) + result = preparer.prepare_inputs( + scheduled_requests=scheduled_requests, + kv_cache_manager=kv_cache_manager, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + enable_spec_decode=True, + runtime_draft_len=1, + is_dummy=False, + ).kwargs + return result + + def _assert_mrope_positions( + self, result: dict[str, Any], rows: list[int] | list[list[int]] + ) -> None: + """`rows` is one list for all three axes, or one list per axis.""" + if not isinstance(rows[0], list): + rows = [rows] * 3 + position_ids = result["position_ids"] + self.assertEqual(tuple(position_ids.shape), (3, 1, len(rows[0]))) + expected = torch.tensor([rows[0:1], rows[1:2], rows[2:3]], dtype=torch.int32, device="cuda") + torch.testing.assert_close(position_ids, expected, atol=0, rtol=0) + + def test_prepare_inputs_spec_decode_mrope_plain_rows_follow_verify_groups(self) -> None: + """Verify groups are packed ahead of plain generation requests, and the + MRoPE delta read slots follow that packing rather than scheduler order. + + The attention kernel indexes the deltas by generation batch index: row + i of the packed generation batch must read the delta of the request + packed i-th, so a verify group's slot comes before every plain + request's, whatever order the scheduler listed them in. + """ + preparer, kv_cache_manager, attn_metadata = self._setup_mrope_inputs( + spec_config=MTPDecodingConfig(max_draft_len=1) + ) + # `verify_*` carry a draft token, `plain_*` do not; `*_mm` have an + # MRoPE delta, `*_text` none. + plain_mm = self._make_mrope_gen_request(6, 1, 0, 20) + verify_mm = self._make_mrope_gen_request(4, 2, 1, 10) + plain_text = self._make_mrope_gen_request(5, 3, 2, None) + verify_text = self._make_mrope_gen_request(3, 4, 3, None) + verify_mm.py_draft_tokens = [7] + verify_text.py_draft_tokens = [7] + + scheduled_requests = ScheduledRequests() + scheduled_requests.context_requests_last_chunk = [] + scheduled_requests.generation_requests = [plain_mm, verify_mm, plain_text, verify_text] + result = self._prepare_mrope_spec_inputs( + preparer, kv_cache_manager, attn_metadata, scheduled_requests + ) + + # Packed order: verify_mm, verify_text, plain_mm, plain_text. The + # verify groups hold two rows each, the plain requests one. + self._assert_mrope_positions(result, [13, 14, 2, 3, 25, 4]) + # Slot 32 is the reserved zero slot (max_num_tokens * pp_size) read by + # the text-only requests. + self.assertEqual(result["mrope_delta_read_seq_slots"].cpu().tolist(), [1, 32, 0, 32]) + self.assertEqual(result["mrope_delta_write_seq_slots"].cpu().tolist(), [1, 0]) + + def test_prepare_inputs_spec_decode_mrope_context_takes_no_read_slot(self) -> None: + """A context request packs ahead of the verify groups with its own + MRoPE positions and takes no generation read slot, so the verify + groups' slots start the list.""" + preparer, kv_cache_manager, attn_metadata = self._setup_mrope_inputs( + spec_config=MTPDecodingConfig(max_draft_len=1) + ) + # A context request with multimodal content takes the placeholder-token + # split, which reads the vocabulary size. + preparer.model.config.vocab_size = 1000 + context = self._make_mrope_ctx_request(5, 1, 0, 30) + verify_mm = self._make_mrope_gen_request(4, 2, 1, 10) + verify_mm.py_draft_tokens = [7] + plain_mm = self._make_mrope_gen_request(6, 3, 2, 20) + + scheduled_requests = ScheduledRequests() + scheduled_requests.context_requests_last_chunk = [context] + scheduled_requests.generation_requests = [plain_mm, verify_mm] + result = self._prepare_mrope_spec_inputs( + preparer, kv_cache_manager, attn_metadata, scheduled_requests + ) + + # The context span keeps the processor's per-axis positions; the + # generation rows after it are text at position + delta on all axes. + self._assert_mrope_positions( + result, + [ + [100, 101, 102, 103, 104, 13, 14, 25], + [1100, 1101, 1102, 1103, 1104, 13, 14, 25], + [2100, 2101, 2102, 2103, 2104, 13, 14, 25], + ], + ) + self.assertEqual(result["mrope_delta_read_seq_slots"].cpu().tolist(), [1, 2]) + # The context request seeds the delta cache first, then the verify + # group and the plain request seed theirs. + self.assertEqual(result["mrope_delta_write_seq_slots"].cpu().tolist(), [0, 1, 2]) + self.assertEqual(context.py_mrope_delta_cache_slot, 0) + + def _setup_extend_as_context_inputs( + self, + ) -> tuple[DecoderInputPreparer, KVCacheManagerV2, AttentionMetadata]: + inputs, manager, metadata = self._setup_mrope_inputs( + spec_config=NGramDecodingConfig(max_draft_len=1), + attention_backend=VanillaAttention, + disable_overlap_scheduler=True, + ) + spec_config = inputs._config.spec_config + self.assertTrue(spec_config.spec_dec_mode.extend_ctx(VanillaAttention)) + self.assertTrue(spec_config.is_linear_tree) + return inputs, manager, metadata + + def test_prepare_inputs_mrope_verify_as_context_takes_no_read_slot(self) -> None: + """When the kernel runs verify groups as chunked context rows, their + rotary comes from `position_ids` alone and only the plain generation + requests take read slots.""" + preparer, kv_cache_manager, attn_metadata = self._setup_extend_as_context_inputs() + verify_mm = self._make_mrope_gen_request(4, 1, 0, 10) + verify_text = self._make_mrope_gen_request(3, 2, 1, None) + plain_mm = self._make_mrope_gen_request(6, 3, 2, 20) + plain_text = self._make_mrope_gen_request(5, 4, 3, None) + verify_mm.py_draft_tokens = [7] + verify_text.py_draft_tokens = [7] + + scheduled_requests = ScheduledRequests() + scheduled_requests.context_requests_last_chunk = [] + scheduled_requests.generation_requests = [verify_mm, verify_text, plain_mm, plain_text] + result = self._prepare_mrope_spec_inputs( + preparer, kv_cache_manager, attn_metadata, scheduled_requests + ) + + self._assert_mrope_positions(result, [13, 14, 2, 3, 25, 4]) + self.assertEqual(result["mrope_delta_read_seq_slots"].cpu().tolist(), [2, 32]) + # The verify group still seeds the delta cache for the plain steps + # that follow it. + self.assertEqual(result["mrope_delta_write_seq_slots"].cpu().tolist(), [0, 2]) + + def test_prepare_inputs_mrope_verify_as_context_alone_emits_no_read_slots(self) -> None: + """A batch whose only delta carriers are verify groups run as context + rows has nothing for the generation-indexed read.""" + preparer, kv_cache_manager, attn_metadata = self._setup_extend_as_context_inputs() + verify_mm = self._make_mrope_gen_request(4, 1, 0, 10) + verify_mm.py_draft_tokens = [7] + plain_text = self._make_mrope_gen_request(5, 2, 1, None) + + scheduled_requests = ScheduledRequests() + scheduled_requests.context_requests_last_chunk = [] + scheduled_requests.generation_requests = [verify_mm, plain_text] + result = self._prepare_mrope_spec_inputs( + preparer, kv_cache_manager, attn_metadata, scheduled_requests + ) + + self._assert_mrope_positions(result, [13, 14, 4]) + self.assertNotIn("mrope_delta_read_seq_slots", result) + self.assertEqual(result["mrope_delta_write_seq_slots"].cpu().tolist(), [0]) + + @skip_pre_blackwell_unittest + def test_ced_warmup_passes_dummy_context_requests(self) -> None: + for warmup in (False, True): + with self.subTest(warmup=warmup): + preparer, manager, metadata = self._setup_mrope_inputs( + 1024, use_mrope=False, disable_overlap_scheduler=True + ) + manager.has_private_swa_suffix = Mock(return_value=True) + preparer.model.model_config.pretrained_config.rope_scaling = None + preparer.model.model = SimpleNamespace( + ced_kv_precompute=True, decoder_replay_split=20 + ) + scheduled = ScheduledRequests() + scheduled.reset_context_requests( + manager.add_dummy_requests([0], token_nums=[257], is_gen=False) + ) + try: + prepared = preparer.prepare_inputs( + scheduled, + manager, + metadata, + enable_spec_decode=False, + runtime_draft_len=0, + is_dummy=warmup, + ) + inputs = prepared.kwargs + self.assertEqual(inputs["context_requests"], scheduled.context_requests) + self.assertTrue(inputs["context_requests"][0].is_dummy) + self.assertEqual(metadata.seq_lens.tolist(), [257]) + self.assertEqual(inputs["input_ids"].numel(), 257) + finally: + for req in scheduled.context_requests: + manager.free_resources(req) + + @skip_pre_blackwell_unittest + def test_ced_inputs_preserve_physical_rows_and_token_sources(self) -> None: + from tensorrt_llm._torch.pyexecutor.ced_replay import EncoderReplay + + for overlap, mixed in ((False, True), (True, True), (True, False)): + with self.subTest(overlap=overlap, mixed=mixed): + preparer, manager, metadata = self._setup_mrope_inputs( + 5001, use_mrope=False, disable_overlap_scheduler=not overlap + ) + manager.has_private_swa_suffix = Mock(return_value=True) + preparer.model.model_config.pretrained_config.rope_scaling = None + preparer.model.model = SimpleNamespace( + ced_kv_precompute=True, decoder_replay_split=20 + ) + scheduled = ScheduledRequests() + if mixed: + context = _create_request_with_tokens(list(range(5000)), 7) + context.context_current_position = 4000 + context.context_chunk_size = 1000 + context.py_seq_slot = 0 + context.py_ced_replay = EncoderReplay(7, 123, 4000, 3872) + scheduled.context_requests_last_chunk = [context] + for index, (length, previous_index) in enumerate(((20, 2), (30, 0))): + req = _create_request(length, 8 + index) + req.py_seq_slot = index + 1 + req.py_batch_idx = previous_index + scheduled.generation_requests.append(req) + new_tokens = torch.tensor([[[222], [999], [111]]], dtype=torch.int32, device="cuda") + prepared = preparer.prepare_inputs( + scheduled, + manager, + metadata, + new_tensors_device=SimpleNamespace(new_tokens=new_tokens) if overlap else None, + enable_spec_decode=False, + runtime_draft_len=0, + is_dummy=False, + ) + inputs = prepared.kwargs + expected_context = list(range(3872, 5000)) if mixed else [] + self.assertEqual( + inputs["input_ids"].cpu().tolist(), + expected_context + ([111, 222] if overlap else [0, 0]), + ) + self.assertEqual( + inputs["position_ids"].view(-1).cpu().tolist(), + expected_context + ([20, 30] if overlap else [19, 29]), + ) + if mixed: + self.assertEqual(inputs["context_requests"], [context]) + self.assertEqual(metadata.seq_lens.tolist(), [1128, 1, 1]) + self.assertEqual(context.context_current_position, 4000) + self.assertEqual(context.context_chunk_size, 1000) + self.assertEqual(context.cached_tokens, 4000) + if not overlap: + self.assertEqual( + list(metadata.kv_cache_params.num_cached_tokens_per_seq), + [3872, 19, 29], + ) + else: + self.assertNotIn("context_requests", inputs) + # GPU sampling is one step ahead of CPU token append. + self.assertEqual( + [r.max_beam_num_tokens for r in scheduled.generation_requests], [20, 30] + ) + + def test_item_scheduled_chunk_embedding_is_not_retained(self) -> None: + """The per-forward item-scheduler payload must not replace the + request's own multimodal data, or each chunked prefill would keep its + last chunk's joined embedding rows until prefill completes.""" + (preparer, kv_cache_manager) = self.enterContext(input_preparer_and_kv_cache()) + preparer.model.config.vocab_size = 100 + attn_metadata = AttentionMetadata( + max_num_requests=4, max_num_tokens=32, kv_cache_manager=kv_cache_manager + ) + attn_metadata.is_cuda_graph = False + + context = _create_request_with_tokens([11, 22, 33, 44], 1) + context.context_chunk_size = 2 + context.py_seq_slot = 0 + context.py_mm_encoder_state = object() + multimodal_data = { + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( + [("image", 0)], output_embedding_lengths=[2] + ) + } + context.py_multimodal_data = multimodal_data + preparer._mm_item_scheduler = Mock() + preparer._mm_item_scheduler.build_multimodal_data_for_llm.return_value = { + "multimodal_embedding": torch.ones((2, 1), dtype=torch.float16), + "multimodal_embedding_is_chunk": True, + } + scheduled_requests = ScheduledRequests() + scheduled_requests.context_requests_chunking = [context] + + preparer.prepare_inputs( + scheduled_requests=scheduled_requests, + kv_cache_manager=kv_cache_manager, + attn_metadata=attn_metadata, + enable_spec_decode=False, + runtime_draft_len=0, + is_dummy=False, + ) + + preparer._mm_item_scheduler.build_multimodal_data_for_llm.assert_called_once() + self.assertIs(context.py_multimodal_data, multimodal_data) + self.assertNotIn("multimodal_embedding", multimodal_data) diff --git a/tests/unittest/_torch/executor/engine/test_decoder_speculative.py b/tests/unittest/_torch/executor/engine/runners/decoder/test_metadata.py similarity index 93% rename from tests/unittest/_torch/executor/engine/test_decoder_speculative.py rename to tests/unittest/_torch/executor/engine/runners/decoder/test_metadata.py index 92e54d10126e..0b3f4e658e38 100644 --- a/tests/unittest/_torch/executor/engine/test_decoder_speculative.py +++ b/tests/unittest/_torch/executor/engine/runners/decoder/test_metadata.py @@ -98,7 +98,9 @@ def test_set_spec_metadata_all_rank_counts_for_one_model() -> None: @pytest.mark.parametrize("dspark", [False, True]) @pytest.mark.parametrize("gather_context_logits", [False, True]) -def test_forward_selects_current_speculative_rows(dspark, gather_context_logits) -> None: +def test_forward_selects_current_speculative_rows( + dspark: bool, gather_context_logits: bool +) -> None: logits = torch.arange(24).reshape(6, 4) metadata = object.__new__(DSparkSpecMetadata) if dspark else SimpleNamespace() runner = object.__new__(DecoderRunner) @@ -106,7 +108,12 @@ def test_forward_selects_current_speculative_rows(dspark, gather_context_logits) runner._inputs = SimpleNamespace(preprocess_inputs=lambda inputs, **kwargs: inputs) runner.sample_in_graph_callable = None - def forward(*, return_context_logits, spec_metadata=None, **kwargs): + def forward( + *, + return_context_logits: bool, + spec_metadata: SimpleNamespace | DSparkSpecMetadata | None = None, + **kwargs: object, + ) -> dict[str, torch.Tensor]: assert return_context_logits is (spec_metadata is not None or gather_context_logits) if dspark and spec_metadata is not None: assert spec_metadata._requires_all_context_logits is gather_context_logits diff --git a/tests/unittest/_torch/executor/engine/runners/decoder/test_runner.py b/tests/unittest/_torch/executor/engine/runners/decoder/test_runner.py new file mode 100644 index 000000000000..681de30a2998 --- /dev/null +++ b/tests/unittest/_torch/executor/engine/runners/decoder/test_runner.py @@ -0,0 +1,912 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + + +from __future__ import annotations + +import contextlib +import unittest +from collections import OrderedDict +from contextlib import nullcontext +from dataclasses import replace +from types import SimpleNamespace +from typing import TYPE_CHECKING +from unittest.mock import Mock, patch + +import pytest +import torch +from _torch.executor.engine.runners.decoder.utils import decoder_config +from _torch.executor.engine_test_utils import _make_request_stub + +import tensorrt_llm._torch.pyexecutor.engine.model_call as model_call_module +import tensorrt_llm._torch.pyexecutor.engine.runners.decoder.runner as decoder_runner_module +from tensorrt_llm._torch.attention.backends.interface import AttentionRuntimeFeatures +from tensorrt_llm._torch.models.checkpoints.hf.weight_mapper import HfWeightMapper +from tensorrt_llm._torch.models.modeling_utils import DecoderModelForCausalLM +from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import KeyType +from tensorrt_llm._torch.pyexecutor.engine.model_call import ModelCaller +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import DecoderRunner +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.inputs import DecoderInputBuffers +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.runner import PrefillCompiledModel +from tensorrt_llm._torch.pyexecutor.engine.runners.interface import PreparedInputs, ScheduledInputs +from tensorrt_llm._torch.pyexecutor.model_loader import ModelLoader +from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManagerType +from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests +from tensorrt_llm._torch.pyexecutor.warmup_timer import _WarmupTimer +from tensorrt_llm._torch.utils import ( + get_model_extra_attrs, + is_torch_compiling, + set_per_request_prefill_cuda_graph_flag, + torch_compiling, +) +from tensorrt_llm.mapping import Mapping + +if TYPE_CHECKING: + from typing import Any + + from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata + from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManager + +_TWO_MODEL_SPEC_MODE = SimpleNamespace(use_one_engine=lambda: False) + + +def _make_decoder_runner( + graph_key: KeyType | None, + runner_enabled: bool = True, + private_decoder: bool = False, + **config_changes: Any, +) -> tuple[DecoderRunner, Mock, Mock, Mock, dict[str, object]]: + config = decoder_config( + **{ + "disable_overlap_scheduler": True, + "is_spec_decode": False, + **config_changes, + } + ) + if config.compile_only_piecewise_graphs: + model = DecoderModelForCausalLM.__new__(DecoderModelForCausalLM) + torch.nn.Module.__init__(model) + model.model = PrefillCompiledModel(torch.nn.Identity(), torch.nn.Identity()) + model.extra_attrs = {} + model.model_config = SimpleNamespace(pretrained_config=SimpleNamespace(rope_scaling=None)) + else: + model = SimpleNamespace( + extra_attrs={}, + model_config=SimpleNamespace(pretrained_config=SimpleNamespace(rope_scaling=None)), + ) + outputs = {"logits": object()} + backend = Mock() + backend.enabled = runner_enabled + backend.pad_batch.side_effect = lambda batch, *_args: nullcontext(batch) + graph_metadata = Mock() + backend.maybe_get_cuda_graph.return_value = ( + (graph_metadata, None, graph_key) if graph_key is not None else (None, None, None) + ) + backend.get_graph_pool.return_value = None + backend.needs_capture.return_value = False + backend.is_warmup_only = False + backend.replay.return_value = outputs + buffers = SimpleNamespace( + input_ids_cuda=torch.zeros(2, dtype=torch.int32), + draft_tokens_cuda=None, + previous_pos_id_offsets_cuda=None, + ) + with ( + patch.object(DecoderInputBuffers, "for_config", return_value=buffers), + patch.object(decoder_runner_module, "CUDAGraphRunner", return_value=backend), + patch.object(decoder_runner_module, "CutlassWorkspaceReclaimer"), + ): + decoder = DecoderRunner( + model, + config, + model_caller=Mock(), + mapping=Mapping(), + dist=None, + moe_load_balancer=None, + sparse_attention_config=None, + attention_runtime_features=AttentionRuntimeFeatures(), + get_runtime_tokens_per_gen_step=Mock(return_value=1), + warmup_timer=_WarmupTimer(rank=0), + iter_states={}, + metrics={}, + ) + metadata = Mock() + decoder._metadata.attn_metadata = metadata + decoder._metadata.setup_attention = Mock(return_value=metadata) + spec_mode = Mock() + spec_mode.attention_need_spec_dec_mode.return_value = False + spec_mode.is_parallel_draft.return_value = False + spec_metadata = Mock( + spec_dec_mode=spec_mode, is_spec_dec_tree=False, is_spec_dec_dynamic_tree=False + ) + decoder._metadata.spec_metadata = spec_metadata + decoder._metadata.setup_speculative = Mock(return_value=spec_metadata) + decoder._inputs.prepare_inputs = Mock( + return_value=PreparedInputs( + { + "prepared": True, + "input_ids": buffers.input_ids_cuda, + } + ) + ) + decoder._inputs.record_promoted_iter_states = Mock() + decoder._inputs.record_input_copy = Mock() + decoder._forward_step = Mock(return_value=outputs) + decoder._execute_logit_post_processors = Mock() + resources = Mock() + peft_cache_manager = Mock(data_type=torch.bfloat16) + + def get_resource_manager(resource_type: ResourceManagerType) -> object: + if resource_type == ResourceManagerType.PEFT_CACHE_MANAGER: + return peft_cache_manager + if resource_type == ResourceManagerType.KV_CACHE_MANAGER: + return SimpleNamespace(has_private_swa_suffix=Mock(return_value=private_decoder)) + return object() + + resources.get_resource_manager.side_effect = get_resource_manager + return decoder, backend, resources, metadata, outputs + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("enabled", [False, True]) +@pytest.mark.parametrize( + "ordinary_capture,eligible", + [(False, True), (False, False), (True, True), (True, False)], +) +def test_decoder_forward_routes_compile_bypass( + enabled: bool, ordinary_capture: bool, eligible: bool +) -> None: + """Exercise both the regular forward and ordinary-graph capture callbacks.""" + key = KeyType(batch_size=1, draft_len=0) if ordinary_capture else None + engine, graph_runner, resources, _, outputs = _make_decoder_runner( + key, compile_only_piecewise_graphs=enabled + ) + eager = torch.nn.Identity() + compiled = torch.nn.Identity() + eager.forward = Mock(return_value=outputs) + compiled.forward = Mock(return_value=outputs) + wrapper = PrefillCompiledModel(eager, compiled) + engine._model.model = wrapper + engine._forward_step.side_effect = lambda *args, **kwargs: wrapper() + graph_runner.needs_capture.return_value = ordinary_capture + graph_runner.is_warmup_only = ordinary_capture + graph_runner.capture.side_effect = lambda key, forward, inputs, **kwargs: forward(inputs) + batch = ScheduledRequests() + batch.generation_requests = [_make_request_stub(1)] + + with ( + patch( + "tensorrt_llm._torch.pyexecutor.engine.runners.decoder.inputs.torch.cuda.Event", + return_value=Mock(), + ), + patch( + "tensorrt_llm._torch.pyexecutor.engine.runners.decoder.runner.get_per_request_prefill_cuda_graph_flag", + return_value=eligible, + ), + ): + result = engine.forward(ScheduledInputs(batch=batch), resource_manager=resources) + + assert result["logits"] is outputs["logits"] + expected_compiled = eligible and not (enabled and ordinary_capture) + selected = compiled.forward if expected_compiled else eager.forward + other = eager.forward if expected_compiled else compiled.forward + selected.assert_called_once_with() + other.assert_not_called() + assert not wrapper._bypass_active + + +@pytest.mark.cpu_only +class DecoderGraphForwardTestCase(unittest.TestCase): + def test_ced_final_context_cannot_promote_to_decode_graph(self) -> None: + key = KeyType(batch_size=1, draft_len=0) + engine, runner, resource_manager, metadata, outputs = _make_decoder_runner( + key, private_decoder=True + ) + engine._model.model = SimpleNamespace(ced_kv_precompute=True, decoder_replay_split=20) + batch = ScheduledRequests() + batch.context_requests_last_chunk = [_make_request_stub(1)] + with ( + patch("tensorrt_llm._torch.pyexecutor.engine.runners.decoder.inputs.torch.cuda.Event"), + patch( + "tensorrt_llm._torch.pyexecutor.engine.runners.decoder.graphs._make_single_token_context_graph_batch" + ) as promote, + ): + actual_outputs = engine.forward( + ScheduledInputs(batch=batch), resource_manager=resource_manager + ) + self.assertEqual(actual_outputs, {**outputs, "runtime_draft_len": 0}) + self.assertIs(actual_outputs["logits"], outputs["logits"]) + promote.assert_not_called() + runner.replay.assert_not_called() + engine._forward_step.assert_called_once() + self.assertIs(engine._inputs.prepare_inputs.call_args.args[2], metadata) + + def test_zero_runtime_draft_speculation_commits_graph_candidate(self) -> None: + key = KeyType(batch_size=2, draft_len=0) + engine, runner, resource_manager, semantic_attn_metadata, outputs = _make_decoder_runner( + key, + spec_config=SimpleNamespace( + draft_len_schedule=None, + spec_dec_mode=_TWO_MODEL_SPEC_MODE, + is_linear_tree=True, + get_runtime_tokens_per_gen_step=Mock(return_value=1), + ), + ) + graph_attn_metadata = runner.maybe_get_cuda_graph.return_value[0] + runner.maybe_get_cuda_graph.return_value = ( + graph_attn_metadata, + engine._metadata.spec_metadata, + key, + ) + context = _make_request_stub(1) + generation = _make_request_stub(2) + batch = ScheduledRequests() + batch.context_requests_last_chunk = [context] + batch.generation_requests = [generation] + + with patch( + "tensorrt_llm._torch.pyexecutor.engine.runners.decoder.inputs.torch.cuda.Event", + return_value=Mock(), + ): + actual_outputs = engine.forward( + ScheduledInputs(batch=batch, enable_spec_decode=True), + resource_manager=resource_manager, + ) + + self.assertEqual(actual_outputs, {**outputs, "runtime_draft_len": 0}) + self.assertIsNot(actual_outputs, outputs) + graph_batch = runner.maybe_get_cuda_graph.call_args.args[0] + self.assertEqual(graph_batch.generation_requests, [context, generation]) + self.assertTrue(runner.maybe_get_cuda_graph.call_args.kwargs["enable_spec_decode"]) + engine._metadata.spec_metadata.update_is_all_greedy_sample.assert_called_once_with( + graph_batch.all_requests() + ) + prepare_args = engine._inputs.prepare_inputs.call_args.args + self.assertIs(prepare_args[0], graph_batch) + self.assertIs(prepare_args[3], engine._metadata.spec_metadata) + self.assertEqual(prepare_args[-1], frozenset({context.py_request_id})) + semantic_attn_metadata.update_spec_dec_param.assert_called_once() + self.assertEqual( + semantic_attn_metadata.update_spec_dec_param.call_args.kwargs["num_contexts"], 1 + ) + prepared_inputs = engine._inputs.prepare_inputs.return_value.kwargs + runner.replay.assert_called_once_with(key, prepared_inputs, position_id_offsets=None) + + def test_zero_runtime_draft_speculation_graph_miss_is_semantic_eager(self) -> None: + engine, runner, resource_manager, semantic_attn_metadata, outputs = _make_decoder_runner( + None, + spec_config=SimpleNamespace( + draft_len_schedule=None, + spec_dec_mode=_TWO_MODEL_SPEC_MODE, + is_linear_tree=True, + get_runtime_tokens_per_gen_step=Mock(return_value=1), + ), + ) + context = _make_request_stub(1) + batch = ScheduledRequests() + batch.context_requests_last_chunk = [context] + + with patch( + "tensorrt_llm._torch.pyexecutor.engine.runners.decoder.inputs.torch.cuda.Event", + return_value=Mock(), + ): + actual_outputs = engine.forward( + ScheduledInputs(batch=batch, enable_spec_decode=True), + resource_manager=resource_manager, + ) + + self.assertEqual(actual_outputs, {**outputs, "runtime_draft_len": 0}) + self.assertIsNot(actual_outputs, outputs) + graph_batch = runner.maybe_get_cuda_graph.call_args.args[0] + self.assertEqual(graph_batch.generation_requests, [context]) + prepare_args = engine._inputs.prepare_inputs.call_args.args + self.assertIs(prepare_args[0], batch) + self.assertIs(prepare_args[2], semantic_attn_metadata) + self.assertIs(prepare_args[3], engine._metadata.spec_metadata) + self.assertEqual(prepare_args[-1], frozenset()) + engine._forward_step.assert_called_once() + runner.replay.assert_not_called() + + def test_zero_runtime_non_linear_tree_speculation_uses_semantic_eager_batch(self) -> None: + engine, runner, resource_manager, semantic_attn_metadata, outputs = _make_decoder_runner( + None, + spec_config=SimpleNamespace( + draft_len_schedule=None, + spec_dec_mode=_TWO_MODEL_SPEC_MODE, + is_linear_tree=False, + tokens_per_gen_step=8, + get_runtime_tokens_per_gen_step=Mock(return_value=1), + ), + ) + context = _make_request_stub(1) + generation = _make_request_stub(2) + batch = ScheduledRequests() + batch.context_requests_last_chunk = [context] + batch.generation_requests = [generation] + + with patch( + "tensorrt_llm._torch.pyexecutor.engine.runners.decoder.inputs.torch.cuda.Event", + return_value=Mock(), + ): + actual_outputs = engine.forward( + ScheduledInputs(batch=batch, enable_spec_decode=True), + resource_manager=resource_manager, + ) + + self.assertEqual(actual_outputs, {**outputs, "runtime_draft_len": 0}) + self.assertIsNot(actual_outputs, outputs) + lookup = runner.maybe_get_cuda_graph.call_args + self.assertIs(lookup.args[0], batch) + self.assertEqual(lookup.kwargs["promoted_context_request_ids"], frozenset()) + prepare_args = engine._inputs.prepare_inputs.call_args.args + self.assertIs(prepare_args[0], batch) + self.assertIs(prepare_args[2], semantic_attn_metadata) + self.assertIs(prepare_args[3], engine._metadata.spec_metadata) + self.assertEqual(prepare_args[-1], frozenset()) + engine._forward_step.assert_called_once() + runner.replay.assert_not_called() + + def test_multimodal_graph_miss_preserves_semantic_payload(self) -> None: + engine, runner, resource_manager, _, _ = _make_decoder_runner(None) + engine._model.config = SimpleNamespace(vocab_size=100) + engine._model.mm_token_ids = torch.tensor([99], dtype=torch.int32) + context = _make_request_stub(1, prompt_len=3) + context.get_tokens = Mock(return_value=[99, 11, 22]) + multimodal_data = { + "multimodal_embedding": object(), + "mrope_config": {"mrope_position_deltas": object()}, + } + context.py_multimodal_data = multimodal_data + batch = ScheduledRequests() + batch.context_requests_last_chunk = [context] + + with patch( + "tensorrt_llm._torch.pyexecutor.engine.runners.decoder.inputs.torch.cuda.Event", + return_value=Mock(), + ): + engine.forward(ScheduledInputs(batch=batch), resource_manager=resource_manager) + + graph_batch = runner.maybe_get_cuda_graph.call_args.args[0] + self.assertIsNot(graph_batch, batch) + self.assertEqual(graph_batch.generation_requests, [context]) + self.assertIs(engine._inputs.prepare_inputs.call_args.args[0], batch) + self.assertIs(context.py_multimodal_data, multimodal_data) + self.assertIn("multimodal_embedding", multimodal_data) + + def test_breakable_graph_falls_back_for_context_logits(self) -> None: + for return_context_logits in (True, False): + with self.subTest(return_context_logits=return_context_logits): + engine, _, resource_manager, _, outputs = _make_decoder_runner(None) + breakable_runner = Mock() + breakable_runner.is_capturing = False + breakable_runner.is_warming_up = False + breakable_runner.has_graph.return_value = True + breakable_runner.execute.return_value = outputs + engine._graphs.breakable_cuda_graph_runner = breakable_runner + prepared = engine._inputs.prepare_inputs.return_value + + def prepare_prefill_graph_batch(*args: object, **kwargs: object) -> PreparedInputs: + set_per_request_prefill_cuda_graph_flag(True) + return prepared + + engine._inputs.prepare_inputs.side_effect = prepare_prefill_graph_batch + self.addCleanup(set_per_request_prefill_cuda_graph_flag, False) + + batch = ScheduledRequests() + batch.context_requests_last_chunk = [_make_request_stub(1)] + batch.context_requests_last_chunk[ + 0 + ].py_return_context_logits = return_context_logits + + with patch( + "tensorrt_llm._torch.pyexecutor.engine.runners.decoder.inputs.torch.cuda.Event", + return_value=Mock(), + ): + actual_outputs = engine.forward( + ScheduledInputs(batch=batch, gather_context_logits=return_context_logits), + resource_manager=resource_manager, + ) + + self.assertEqual(actual_outputs, {**outputs, "runtime_draft_len": 0}) + self.assertIsNot(actual_outputs, outputs) + if return_context_logits: + breakable_runner.execute.assert_not_called() + engine._forward_step.assert_called_once() + else: + breakable_runner.execute.assert_called_once() + engine._forward_step.assert_not_called() + + def test_generation_only_forward_does_not_call_new_selector(self) -> None: + key = KeyType(batch_size=1, draft_len=0) + engine, runner, resource_manager, _, _ = _make_decoder_runner(key) + generation = _make_request_stub(2) + batch = ScheduledRequests() + batch.generation_requests = [generation] + + with patch( + "tensorrt_llm._torch.pyexecutor.engine.runners.decoder.inputs.torch.cuda.Event", + return_value=Mock(), + ): + engine.forward(ScheduledInputs(batch=batch), resource_manager=resource_manager) + + lookup = runner.maybe_get_cuda_graph.call_args + self.assertIs(lookup.args[0], batch) + self.assertEqual(lookup.kwargs["promoted_context_request_ids"], frozenset()) + self.assertFalse(lookup.kwargs["use_lora_graph"]) + self.assertIs(engine._inputs.prepare_inputs.call_args.args[0], batch) + self.assertEqual(engine._inputs.prepare_inputs.call_args.args[-1], frozenset()) + + def test_generation_lora_request_selects_lora_graph(self) -> None: + key = KeyType(batch_size=1, draft_len=0, use_lora_graph=True) + engine, runner, resource_manager, _, _ = _make_decoder_runner( + key, cuda_graph_specialize_lora=True + ) + engine._lora = engine._graphs._lora = SimpleNamespace(cuda_graph_manager=object()) + generation = _make_request_stub(2) + generation.lora_task_id = 42 + batch = ScheduledRequests() + batch.generation_requests = [generation] + + with patch( + "tensorrt_llm._torch.pyexecutor.engine.runners.decoder.inputs.torch.cuda.Event", + return_value=Mock(), + ): + engine.forward(ScheduledInputs(batch=batch), resource_manager=resource_manager) + + self.assertTrue(runner.maybe_get_cuda_graph.call_args.kwargs["use_lora_graph"]) + self.assertEqual( + runner.maybe_get_cuda_graph.call_args.kwargs["peft_cache_data_type"], torch.bfloat16 + ) + self.assertTrue(engine._inputs.prepare_inputs.call_args.kwargs["use_lora_graph"]) + + def test_global_incompatibilities_bypass_candidate_selection(self) -> None: + cases = ( + "graphs_disabled", + "speculative_nonzero_runtime_draft", + "beam", + "extra_inputs", + "ple_recurrent_state", + "nested_ple_recurrent_state", + "context_parallel", + "context_only_cache", + ) + for case in cases: + with self.subTest(case=case): + engine, runner, resource_manager, _, _ = _make_decoder_runner( + None, + max_beam_width=2 if case == "beam" else 1, + spec_config=SimpleNamespace( + draft_len_schedule=None, + spec_dec_mode=_TWO_MODEL_SPEC_MODE, + is_linear_tree=True, + get_runtime_tokens_per_gen_step=Mock(return_value=1), + ) + if case == "speculative_nonzero_runtime_draft" + else None, + ) + spec_state = {} + if case == "graphs_disabled": + runner.enabled = False + elif case == "speculative_nonzero_runtime_draft": + spec_state = dict(enable_spec_decode=True, runtime_draft_len=1) + elif case == "extra_inputs": + engine._graphs._context_promotion_supported = False + elif case == "ple_recurrent_state": + engine._model.has_ple = True + elif case == "nested_ple_recurrent_state": + engine._model.model = SimpleNamespace( + llm=SimpleNamespace(model=SimpleNamespace(has_ple=True)) + ) + elif case == "context_parallel": + engine._mapping.cp_size = 2 + elif case == "context_only_cache": + resource_manager.get_resource_manager.side_effect = lambda _: SimpleNamespace( + context_swa_layer_limit=20 + ) + + batch = ScheduledRequests() + batch.context_requests_last_chunk = [_make_request_stub(1)] + with patch( + "tensorrt_llm._torch.pyexecutor.engine.runners.decoder.inputs.torch.cuda.Event", + return_value=Mock(), + ): + engine.forward( + ScheduledInputs(batch=batch, **spec_state), + resource_manager=resource_manager, + ) + + lookup = runner.maybe_get_cuda_graph.call_args + self.assertIs(lookup.args[0], batch) + self.assertEqual(lookup.kwargs["promoted_context_request_ids"], frozenset()) + self.assertIs(engine._inputs.prepare_inputs.call_args.args[0], batch) + + def test_forward_threads_call_state_and_returns_prepared_length(self) -> None: + engine, runner, resource_manager, _, outputs = _make_decoder_runner( + None, + spec_config=SimpleNamespace( + draft_len_schedule=None, + spec_dec_mode=_TWO_MODEL_SPEC_MODE, + is_linear_tree=False, + tokens_per_gen_step=8, + get_runtime_tokens_per_gen_step=Mock(return_value=1), + ), + ) + # A non-linear tree widens the length to its total draft tokens during + # preparation, after graph selection. + prepared = engine._inputs.prepare_inputs.return_value + engine._inputs.prepare_inputs.return_value = replace(prepared, runtime_draft_len=7) + batch = ScheduledRequests() + batch.generation_requests = [_make_request_stub(2)] + + with patch( + "tensorrt_llm._torch.pyexecutor.engine.runners.decoder.inputs.torch.cuda.Event", + return_value=Mock(), + ): + actual_outputs = engine.forward( + ScheduledInputs(batch=batch, enable_spec_decode=True, runtime_draft_len=2), + resource_manager=resource_manager, + ) + + self.assertEqual(actual_outputs, {**outputs, "runtime_draft_len": 7}) + self.assertEqual(runner.pad_batch.call_args.args[2], 2) + self.assertTrue(runner.maybe_get_cuda_graph.call_args.kwargs["enable_spec_decode"]) + prepare_kwargs = engine._inputs.prepare_inputs.call_args.kwargs + self.assertEqual( + ( + prepare_kwargs["enable_spec_decode"], + prepare_kwargs["runtime_draft_len"], + prepare_kwargs["is_dummy"], + ), + (True, 2, False), + ) + step_kwargs = engine._forward_step.call_args.kwargs + self.assertEqual( + ( + step_kwargs["enable_spec_decode"], + step_kwargs["runtime_draft_len"], + step_kwargs["is_dummy"], + ), + (True, 7, False), + ) + + +@pytest.mark.cpu_only +class TestDecodeGraphMetadataPreparation(unittest.TestCase): + @staticmethod + def _metadata() -> CSA2TrtllmMetadata: + from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata + + metadata = object.__new__(CSA2TrtllmMetadata) + metadata.is_cuda_graph = True + return metadata + + def test_scope_requires_selected_ordinary_decode_graph(self) -> None: + key = KeyType(batch_size=1, draft_len=0) + cases = ( + "ordinary", + "ordinary_overlap", + "graph_miss", + "eager_metadata", + "context", + "promoted_context", + "speculative", + "beam", + "encoder_decoder", + "context_parallel", + "draft_key", + ) + for case in cases: + with self.subTest(case=case): + engine, _, _, _, _ = _make_decoder_runner( + key, + disable_overlap_scheduler=case != "ordinary_overlap", + max_beam_width=2 if case == "beam" else 1, + ) + enable_spec_decode = False + metadata = self._metadata() + batch = ScheduledRequests() + batch.generation_requests = [_make_request_stub(1)] + selected = key + promoted = frozenset() + if case == "graph_miss": + selected = None + elif case == "eager_metadata": + metadata.is_cuda_graph = False + elif case == "context": + batch.context_requests_last_chunk = [_make_request_stub(2)] + elif case == "promoted_context": + promoted = frozenset({1}) + elif case == "speculative": + enable_spec_decode = True + elif case == "encoder_decoder": + engine._deferred_decode_outputs_supported = False + elif case == "context_parallel": + engine._mapping.cp_size = 2 + elif case == "draft_key": + selected = key._replace(draft_len=1) + with engine._cuda_graph_metadata_prepare_scope( + metadata, selected, batch, promoted, enable_spec_decode=enable_spec_decode + ): + self.assertEqual( + metadata._csa2_defer_decode_outputs, + case in ("ordinary", "ordinary_overlap"), + ) + self.assertFalse(metadata._csa2_defer_decode_outputs) + + def test_forward_scopes_proof_to_prepare_for_capture_replay_and_fallback(self) -> None: + key = KeyType(batch_size=1, draft_len=0) + for mode in ("warmup", "capture", "replay", "fallback", "prepare_error"): + with self.subTest(mode=mode): + selected = None if mode == "fallback" else key + engine, runner, resource_manager, _, outputs = _make_decoder_runner(selected) + metadata = self._metadata() + engine.attn_metadata = metadata + engine._metadata.setup_attention.return_value = metadata + runner.maybe_get_cuda_graph.return_value = (metadata, None, selected) + runner.needs_capture.return_value = mode in ("warmup", "capture") + runner.is_warmup_only = mode == "warmup" + runner.capture.return_value = outputs + batch = ScheduledRequests() + batch.generation_requests = [_make_request_stub(1)] + prepared = engine._inputs.prepare_inputs.return_value + + def prepare(*args: object, **kwargs: object) -> PreparedInputs: + self.assertIs(metadata.trtllm_gen_jit_warmup, False) + self.assertEqual(metadata._csa2_defer_decode_outputs, mode != "fallback") + if mode == "prepare_error": + raise RuntimeError("prepare failed") + return prepared + + engine._inputs.prepare_inputs.side_effect = prepare + with patch( + "tensorrt_llm._torch.pyexecutor.engine.runners.decoder.inputs.torch.cuda.Event", + return_value=Mock(), + ): + if mode == "prepare_error": + with self.assertRaisesRegex(RuntimeError, "prepare failed"): + engine.forward( + ScheduledInputs(batch=batch), resource_manager=resource_manager + ) + else: + actual_outputs = engine.forward( + ScheduledInputs(batch=batch), resource_manager=resource_manager + ) + self.assertEqual(actual_outputs, {**outputs, "runtime_draft_len": 0}) + self.assertIs(actual_outputs["logits"], outputs["logits"]) + self.assertFalse(metadata._csa2_defer_decode_outputs) + if mode in ("warmup", "capture"): + runner.capture.assert_called_once() + if mode in ("capture", "replay"): + runner.replay.assert_called_once() + if mode == "fallback": + engine._forward_step.assert_called_once() + + def test_jit_warmup_flag_marks_only_the_warmup_forward(self) -> None: + engine, _, resource_manager, _, _ = _make_decoder_runner(None) + runner = engine + metadata = self._metadata() + runner.attn_metadata = metadata + runner._metadata.setup_attention.return_value = metadata + prepared = runner._inputs.prepare_inputs.return_value + observed = [] + + def prepare(*args: object, **kwargs: object) -> PreparedInputs: + observed.append(metadata.trtllm_gen_jit_warmup) + return prepared + + runner._inputs.prepare_inputs.side_effect = prepare + batch = ScheduledRequests() + batch.generation_requests = [_make_request_stub(1)] + for jit_warmup in (True, False): + runner._forward_warmup( + batch, + resource_manager, + enable_spec_decode=False, + runtime_draft_len=0, + trtllm_gen_jit_warmup=jit_warmup, + ) + self.assertEqual(observed, [True, False]) + self.assertFalse(runner._trtllm_gen_jit_warmup) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("local_contexts", [0, 1]) +def test_prefill_compile_uses_all_rank_prefill_decision( + monkeypatch: pytest.MonkeyPatch, + local_contexts: int, +) -> None: + """Route using the global prefill decision, regardless of local contexts.""" + eager = torch.nn.Linear(4, 4) + compiled = torch.nn.Module() + compiled.shared = eager + expected = torch.ones((2, 4)) + eager.forward = Mock(return_value=expected) + compiled.forward = Mock(return_value=expected) + router = PrefillCompiledModel(eager, compiled) + assert router.weight is eager.weight + assert list(router.parameters()) == list(eager.parameters()) + # Local decode-only ranks participate when another attention-DP rank + # prefills; an over-ceiling local context batch uses eager instead. + metadata = SimpleNamespace(num_contexts=local_contexts) + for eligible in (True, False): + monkeypatch.setattr( + decoder_runner_module, "get_per_request_prefill_cuda_graph_flag", lambda: eligible + ) + assert router(expected, attn_metadata=metadata) is expected + selected = compiled if eligible else eager + selected.forward.assert_called_once_with(expected, attn_metadata=metadata) + assert compiled.forward.call_count == eager.forward.call_count == 1 + + +@pytest.mark.cpu_only +def test_prefill_compile_preserves_partial_weight_reload( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Reload selected weights by original names into both model entry points.""" + import tensorrt_llm._torch.models.modeling_utils as modeling_utils + + eager = torch.nn.Sequential(OrderedDict(layers=torch.nn.Sequential(torch.nn.Linear(4, 4)))) + compiled = torch.compile(eager, backend="eager") + model = DecoderModelForCausalLM.__new__(DecoderModelForCausalLM) + torch.nn.Module.__init__(model) + model.model_config = SimpleNamespace( + pretrained_config=SimpleNamespace(tie_word_embeddings=False) + ) + model.model = PrefillCompiledModel(eager, compiled) + mapper = HfWeightMapper() + mapper._model = model + loader = ModelLoader.__new__(ModelLoader) + loader.weight_mapper = mapper + monkeypatch.setenv("TRT_LLM_DISABLE_LOAD_WEIGHTS_IN_PARALLEL", "True") + monkeypatch.setattr(modeling_utils, "local_mpi_rank", lambda: 0) + monkeypatch.setattr(torch.cuda, "set_device", lambda device: None) + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + + weight = eager.layers[0].weight + old_bias = eager.layers[0].bias.detach().clone() + replacement = torch.full_like(weight, 3) + for remove_duplicate in (False, True): + names = dict(model.named_modules(remove_duplicate=remove_duplicate)) + assert names["model.layers.0"] is eager.layers[0] + assert not any("eager_model" in name or "compiled_model" in name for name in names) + assert list(model.model.children()) == [eager] + assert list(model.model.state_dict()) == [ + "eager_model.layers.0.weight", + "eager_model.layers.0.bias", + ] + apply_tensor = Mock(side_effect=lambda tensor: tensor) + model.model._apply(apply_tensor) + assert apply_tensor.call_count == 2 + loader.reload(model, {"model.layers.0.weight": replacement}, allow_partial_loading=True) + + assert eager.layers[0].weight is weight + assert compiled.layers[0].weight is weight + torch.testing.assert_close(weight, replacement) + torch.testing.assert_close(eager.layers[0].bias, old_bias) + inputs = torch.ones(2, 4) + for eligible in (False, True): + monkeypatch.setattr( + decoder_runner_module, "get_per_request_prefill_cuda_graph_flag", lambda: eligible + ) + torch.testing.assert_close(model.model(inputs), inputs @ replacement.t() + old_bias) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("prefill_only", [False, True]) +@pytest.mark.parametrize("eligible", [False, True]) +@pytest.mark.parametrize("bypass", [False, True]) +@pytest.mark.parametrize("raises", [False, True]) +def test_prefill_compile_scopes_whole_model_forward( + monkeypatch: pytest.MonkeyPatch, + prefill_only: bool, + eligible: bool, + bypass: bool, + raises: bool, +) -> None: + """Scope compile dispatch by phase independently of the model wrapper.""" + observed = [] + + def forward(**kwargs: object) -> str: + """Record compile state as a stand-in for a model-specific epilogue.""" + observed.append(is_torch_compiling()) + # This represents work after the transformer, such as Eagle3 drafting. + if raises: + raise RuntimeError("epilogue failure") + observed.append(is_torch_compiling()) + return "done" + + eager_model = torch.nn.Identity() + eager_model.model_config = SimpleNamespace(extra_attrs={}) + eager_model.forward = forward + compiled_model = torch.nn.Identity() + compiled_model.forward = forward + prefill_compiled_model = PrefillCompiledModel(eager_model, compiled_model) + model = prefill_compiled_model if prefill_only else eager_model + engine = SimpleNamespace( + _model_caller=ModelCaller(model, prefill_compile_only=prefill_only), + _eager_workspace_reclaimer=None, + _moe_workspace_reclaimer=None, + ) + monkeypatch.setattr(model_call_module, "get_model_extra_attrs", lambda: {}) + monkeypatch.setattr( + model_call_module, "get_per_request_prefill_cuda_graph_flag", lambda: eligible + ) + monkeypatch.setattr(model_call_module, "is_trace_enabled", lambda name: False) + bypass_scope = ( + prefill_compiled_model.bypass() if prefill_only and bypass else contextlib.nullcontext() + ) + with torch_compiling(True), bypass_scope: + if raises: + with pytest.raises(RuntimeError, match="epilogue failure"): + DecoderRunner._model_forward(engine, is_dummy=False, attn_metadata=Mock()) + else: + assert ( + DecoderRunner._model_forward(engine, is_dummy=False, attn_metadata=Mock()) == "done" + ) + assert is_torch_compiling() + expected = eligible if prefill_only else True + assert observed == [expected] * (1 if raises else 2) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("is_dummy", [False, True]) +def test_decoder_runner_passes_call_state_without_engine_writes(is_dummy: bool) -> None: + engine, _, resources, _, _ = _make_decoder_runner(None) + # The dummy flag comes from the caller, not from the warmup scope. + inputs = ScheduledInputs( + batch=ScheduledRequests(), enable_spec_decode=True, runtime_draft_len=2 + ) + expected = {"logits": object(), "runtime_draft_len": 4} + engine._forward_decoder = Mock(return_value=expected) + + outputs = engine.forward(inputs, resource_manager=resources, is_dummy=is_dummy) + + assert outputs is expected + engine._forward_decoder.assert_called_once_with(inputs, resources, is_dummy=is_dummy) + assert (inputs.enable_spec_decode, inputs.runtime_draft_len) == (True, 2) + + +@pytest.mark.cpu_only +def test_forward_warmup_runs_local_call_state_in_forward_context() -> None: + engine, _, resources, _, _ = _make_decoder_runner(None) + batch = ScheduledRequests() + + def decoder_forward( + inputs: ScheduledInputs, resource_manager: ResourceManager, *, is_dummy: bool + ) -> dict[str, object]: + assert torch.is_inference_mode_enabled() + assert get_model_extra_attrs() is engine._model.extra_attrs + return {"logits": object(), "runtime_draft_len": 3} + + engine._forward_decoder = Mock(side_effect=decoder_forward) + + outputs = engine._forward_warmup(batch, resources, enable_spec_decode=True, runtime_draft_len=1) + + inputs = engine._forward_decoder.call_args.args[0] + assert inputs.batch is batch + assert inputs.new_tensors_device is None + assert inputs.cache_indirection_buffer is None + assert (inputs.enable_spec_decode, inputs.runtime_draft_len) == (True, 1) + assert engine._forward_decoder.call_args.kwargs == {"is_dummy": True} + assert outputs["runtime_draft_len"] == 3 + assert get_model_extra_attrs() is None + + +@pytest.mark.cpu_only +def test_decoder_runner_forward_enters_model_context() -> None: + engine, _, resources, _, _ = _make_decoder_runner(None) + inputs = ScheduledInputs(batch=ScheduledRequests()) + + def decoder_forward( + inputs: ScheduledInputs, resource_manager: ResourceManager, *, is_dummy: bool + ) -> dict[str, object]: + assert torch.is_inference_mode_enabled() + assert get_model_extra_attrs() is engine._model.extra_attrs + return {"logits": object()} + + engine._forward_decoder = Mock(side_effect=decoder_forward) + + engine.forward(inputs, resource_manager=resources) + + engine._forward_decoder.assert_called_once_with(inputs, resources, is_dummy=False) + assert get_model_extra_attrs() is None diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py b/tests/unittest/_torch/executor/engine/runners/decoder/test_warmup.py similarity index 53% rename from tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py rename to tests/unittest/_torch/executor/engine/runners/decoder/test_warmup.py index ee9db7c8a56d..2ca5fbd39a61 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py +++ b/tests/unittest/_torch/executor/engine/runners/decoder/test_warmup.py @@ -1,341 +1,333 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Unit tests for warmup-cleanup behavior in ModelEngine.warmup(). -Locks in that gc.collect() + torch.cuda.empty_cache() fire immediately after -_run_autotuner_warmup (step b) to release autotuner exploration leftovers. - -The torch.cuda.empty_cache() after teardown_managers() in py_executor_creator -is covered end-to-end by integration tests rather than unit-tested here. -""" +from __future__ import annotations import contextlib import gc import os import sys import unittest -from collections import OrderedDict from collections.abc import Iterator -from dataclasses import dataclass, replace +from dataclasses import replace from types import ModuleType, SimpleNamespace +from typing import TYPE_CHECKING +from unittest import mock from unittest.mock import Mock, call, patch import pytest import torch +from _torch.executor.engine.runners.decoder.utils import ( + decoder_config, + decoder_graphs, + decoder_warmup, +) +from _torch.executor.engine_test_utils import ( + DummyModel, + _agreement, + _capture_tllm_logs, + kv_cache_manager, +) -import tensorrt_llm import tensorrt_llm._torch.pyexecutor.engine.model_call as model_call_module -import tensorrt_llm._torch.pyexecutor.engine.runners.decoder.runner as decoder_runner_module import tensorrt_llm._torch.pyexecutor.engine.runners.decoder.warmup as decoder_warmup_module -import tensorrt_llm._torch.pyexecutor.model_engine as model_engine_module +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.custom_ops.torch_custom_ops import MXFP8GemmRunner -from tensorrt_llm._torch.model_config import ModelConfig -from tensorrt_llm._torch.models.checkpoints.hf.weight_mapper import HfWeightMapper -from tensorrt_llm._torch.models.modeling_utils import DecoderModelForCausalLM, remove_weights +from tensorrt_llm._torch.models.modeling_utils import remove_weights from tensorrt_llm._torch.modules.linear import MXFP8LinearMethod from tensorrt_llm._torch.pyexecutor.engine.model_call import ModelCaller from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import DecoderRunner -from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.graphs import DecoderCudaGraphs -from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.runner import PrefillCompiledModel +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import warmup as warmup_module +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.metadata import DecoderMetadata from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.warmup import DecoderWarmup -from tensorrt_llm._torch.pyexecutor.engine.runners.encoder_decoder import EncoderDecoderRunner -from tensorrt_llm._torch.pyexecutor.engine.runners.no_kv_cache import NoKVCacheRunner from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 -from tensorrt_llm._torch.pyexecutor.model_engine import ModelEngine -from tensorrt_llm._torch.pyexecutor.model_loader import ModelLoader from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManager, ResourceManagerType from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm._torch.pyexecutor.warmup_timer import _WarmupTimer -from tensorrt_llm._torch.speculative.utils import resolve_draft_len, update_draft_len from tensorrt_llm._torch.utils import is_torch_compiling, torch_compiling -from tensorrt_llm.llmapi import CudaGraphConfig, KvCacheConfig -from tensorrt_llm.llmapi.llm_args import ( - DecodingBaseConfig, - DraftTargetDecodingConfig, - MultimodalConfig, - PARDDecodingConfig, - PrefillCudaGraphBackend, - TorchLlmArgs, -) -from tensorrt_llm.mapping import Mapping +from tensorrt_llm.llmapi.llm_args import DraftTargetDecodingConfig, PrefillCudaGraphBackend +if TYPE_CHECKING: + from collections.abc import Callable -@pytest.mark.cpu_only -@pytest.mark.parametrize("model_kind", ["default", "m3", "m3_vl", "gdn", "mamba"]) -@pytest.mark.parametrize("compile_enabled,piecewise", [(False, False), (True, False), (True, True)]) -@pytest.mark.parametrize("compile_only_piecewise_graphs", [False, True]) -def test_pcg_fx_fallback_policy_is_model_specific( - monkeypatch: pytest.MonkeyPatch, - model_kind: str, - compile_enabled: bool, - piecewise: bool, - compile_only_piecewise_graphs: bool, -) -> None: - """Exercise the real constructor's routing gate without loading weights or CUDA.""" - from tensorrt_llm._torch.models.modeling_minimaxm3 import ( - MiniMaxM3ForCausalLM, - MiniMaxM3VLForConditionalGeneration, + from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.config import DecoderRunnerConfig + + +@contextlib.contextmanager +def _warmup_and_resources() -> Iterator[tuple[DecoderWarmup, ResourceManager]]: + config = decoder_config() + model = DummyModel(torch.half) + backend = Mock(enabled=True, is_warmup_only=False) + backend.allow_capture.side_effect = contextlib.nullcontext + graphs = decoder_graphs(config, model=model, cuda_graph_runner=backend) + warmup = decoder_warmup(config, model=model, graphs=graphs) + with kv_cache_manager(model=model, max_num_tokens=config.max_num_tokens) as manager: + yield warmup, ResourceManager({ResourceManagerType.KV_CACHE_MANAGER: manager}) + + +class _Tracker: + """Records method-call order (and each call's arguments) via mock side_effects.""" + + def __init__(self) -> None: + self.calls: list[str] = [] + self.call_args: dict[str, list[tuple[tuple[object, ...], dict[str, object]]]] = {} + + def __call__(self, name: str) -> Callable[..., None]: + def _wrapped(*args: object, **kwargs: object) -> None: + self.calls.append(name) + self.call_args.setdefault(name, []).append((args, kwargs)) + + return _wrapped + + +def _run_warmup_tracked( + model_engine: DecoderWarmup, + resource_manager: ResourceManager, + *, + force_helix_cp: bool = False, + capture_logs: bool = False, + tracker: _Tracker | None = None, +) -> tuple[list[str], list[str] | None]: + """Patch warmup stages and cleanup operations, then run warmup. + + Optionally force helix CP and capture logs. Pass ``tracker`` to inspect the + recorded call arguments. Returns (call_order_list, log_records_or_None). + """ + if tracker is None: + tracker = _Tracker() + # Engines left by earlier tests are cyclic; finalize them before tracking. + gc.collect() + # Observe this module's explicit cleanup without counting unrelated + # destructors triggered by gc.collect() or lazy imports. + tracked_cuda = SimpleNamespace(**vars(torch.cuda)) + tracked_cuda.empty_cache = tracker("empty_cache") + tracked_torch = SimpleNamespace(**vars(torch)) + tracked_torch.cuda = tracked_cuda + helix_ctx = ( + patch.object(model_engine.mapping, "has_cp_helix", return_value=True) + if force_helix_cp + else contextlib.nullcontext() ) - from tensorrt_llm._torch.models.modeling_nemotron_h import NemotronHForCausalLM - from tensorrt_llm._torch.models.modeling_qwen3_next import Qwen3NextForCausalLM - from tensorrt_llm._torch.pyexecutor import _util - - model_cls = { - "default": DecoderModelForCausalLM, - "m3": MiniMaxM3ForCausalLM, - "m3_vl": MiniMaxM3VLForConditionalGeneration, - "gdn": Qwen3NextForCausalLM, - "mamba": NemotronHForCausalLM, - }[model_kind] - assert model_cls.use_fx_for_pcg_fallback is (model_kind not in ("m3", "m3_vl")) - model = model_cls.__new__(model_cls) - torch.nn.Module.__init__(model) - mapping = Mapping() - model.model_config = SimpleNamespace( - pretrained_config=SimpleNamespace(torch_dtype=torch.float32, hidden_size=4), - mapping=mapping, - sparse_attention_config=None, + + with ( + helix_ctx, + patch.object(model_engine, "_general_warmup", side_effect=tracker("general_warmup")), + patch.object(model_engine, "_run_autotuner_warmup", side_effect=tracker("autotuner")), + patch.object( + model_engine._graphs, + "_capture_generation_cuda_graphs", + side_effect=tracker("generation_cuda_graph"), + ), + patch.object( + model_engine._graphs, + "_capture_additional_cuda_graphs", + side_effect=tracker("additional_cuda_graph"), + ), + patch.object( + model_engine._graphs, + "_capture_prefill_cuda_graphs", + side_effect=tracker("prefill_cuda_graph"), + ), + patch( + "tensorrt_llm._torch.pyexecutor.engine.runners.decoder.warmup.warmup_sampling_module", + side_effect=tracker("sampling_warmup"), + ), + patch.object(decoder_warmup_module, "torch", tracked_torch), + patch( + "tensorrt_llm._torch.custom_ops.torch_custom_ops.MoERunner.clear_all_workspaces", + side_effect=tracker("moe_clear"), + ), + ): + if capture_logs: + with _capture_tllm_logs() as logs: + model_engine.run(resource_manager) + return tracker.calls, logs + model_engine.run(resource_manager) + return tracker.calls, None + + +def _indexer_q_prewarm_owner( + modules: list[torch.nn.Module | SimpleNamespace | Mock], tp_size: int +) -> SimpleNamespace: + return SimpleNamespace( + model=SimpleNamespace(modules=lambda: iter(modules)), + mapping=SimpleNamespace(tp_size=tp_size), + dist=Mock(), ) - eager = torch.nn.Linear(4, 4) - model.model = eager - compile_config = SimpleNamespace( - enable_fullgraph=True, - compile_only_piecewise_graphs=compile_only_piecewise_graphs, - enable_inductor=False, - enable_userbuffers=False, - max_num_streams=1, + + +def _deepseek_v4_indexer( + q_lora_rank: int = 256, with_wq_b: bool = True, fused_enabled: bool = True +) -> Mock: + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4 import DeepseekV4Indexer + + indexer = Mock(spec=DeepseekV4Indexer) + indexer.wq_b = ( + SimpleNamespace(weight=torch.empty((64, q_lora_rank), dtype=torch.bfloat16)) + if with_wq_b + else None ) - llm_args = SimpleNamespace( - encode_only=False, - mm_encoder_only=False, - get_runtime_sizes=lambda: (1, 16, 64, 4), - encoder_max_batch_size=None, - encoder_max_num_tokens=None, - enable_in_graph_sampling=False, - enable_return_routed_experts=False, - multimodal_config=MultimodalConfig(), - checkpoint_format="HF", - trust_remote_code=False, - disable_overlap_scheduler=True, - kv_cache_config=KvCacheConfig(), - sleep_config=None, - enable_layerwise_nvtx_marker=False, - cuda_graph_config=None, - torch_compile_config=compile_config if compile_enabled else None, - prefill_cuda_graph_backend=PrefillCudaGraphBackend.PIECEWISE if piecewise else None, - prefill_capture_num_tokens=[4], - allreduce_strategy="AUTO", - attn_backend="TRTLLM", - sparse_attention_config=None, + indexer._is_fused_project_mxfp4_enabled.return_value = fused_enabled + return indexer + + +def _other_pp_stage_indexer() -> Mock: + """Layers owned by another pipeline stage keep wq_b but lose its weight.""" + indexer = _deepseek_v4_indexer() + indexer.wq_b = torch.nn.Linear(256, 64, bias=False) + remove_weights(indexer.wq_b) + return indexer + + +class _StandInMambaCacheManager: + """Concrete stand-in for the abstract Mamba cache-manager type check.""" + + def __init__(self, available_tokens: int = 8) -> None: + self._available_tokens = available_tokens + + def get_num_available_tokens(self, **_kwargs: object) -> int: + return self._available_tokens + + +def _engine( + *, world_size: int = 1, dwdp_size: int = 0, config: DecoderRunnerConfig | None = None +) -> DecoderWarmup: + agreement = _agreement(world_size=world_size, dwdp_size=dwdp_size) + engine = decoder_warmup( + config or decoder_config(), agreement=agreement, mapping=agreement.mapping ) - for name in ( - "should_enable_overlap_headroom", - "resolved_kv_cache_manager_is_v2", - ): - monkeypatch.setattr(_util, name, Mock(return_value=False)) - monkeypatch.setattr(_util, "compute_max_num_sequences", Mock(return_value=4)) - for name in ( - "_configure_deep_gemm_pdl", - "create_input_processor", - "setup_mm_encoder_attn_metadata", + engine._reset_moe_alltoall_state = mock.Mock() + return engine + + +@contextlib.contextmanager +def _released_batch(batch: object) -> Iterator[object]: + yield batch + + +@contextlib.contextmanager +def _no_cuda_side_effects() -> Iterator[tuple[mock.Mock, mock.Mock]]: + with ( + mock.patch.object(torch.cuda, "empty_cache") as empty_cache, + mock.patch.object(torch.cuda, "synchronize") as synchronize, ): - monkeypatch.setattr(model_engine_module, name, Mock()) - monkeypatch.setattr( - model_engine_module, "resolve_mrope_position_deltas_cache", lambda model: None - ) - monkeypatch.setattr( - model_engine_module, "is_hybrid_linear", lambda config: model_kind in ("gdn", "mamba") + yield empty_cache, synchronize + + +def _general_warmup_engine(*, world_size: int, dwdp_size: int) -> DecoderWarmup: + engine = _engine(world_size=world_size, dwdp_size=dwdp_size) + engine._warmup_timer = _WarmupTimer(rank=engine.dist.rank) + batch = object() + engine._dummy_requests = SimpleNamespace( + create_warmup_request=mock.Mock(return_value=batch), + release_batch_context=lambda *_a, **_kw: _released_batch(batch), ) - monkeypatch.setattr( - model_engine_module.MultimodalItemScheduler, "maybe_create", Mock(return_value=None) + return engine + + +_KV_ALLOC_ERROR = "Can't allocate new blocks for window size 8" + + +def _mamba_engine(*, world_size: int = 1, dwdp_size: int = 0) -> tuple[DecoderWarmup, object]: + engine = _engine( + world_size=world_size, + dwdp_size=dwdp_size, + config=decoder_config( + max_num_tokens=8, max_batch_size=4, max_seq_len=8, enable_autotuner=False + ), ) - monkeypatch.setattr(ModelEngine, "_validate_breakable_cuda_graph_compatibility", Mock()) - monkeypatch.setattr(ModelEngine, "_init_model_capacity", Mock()) - monkeypatch.setattr(ModelEngine, "__del__", lambda self: None) - backend_factory = Mock() - backend_factory.Streams = list - monkeypatch.setattr(model_engine_module, "Backend", backend_factory) - compiled = torch.nn.Identity() - compile_model = Mock(return_value=compiled) - monkeypatch.setattr(torch, "compile", compile_model) - monkeypatch.setattr(torch._dynamo.config, "cache_size_limit", 16) - # Stop after the complete compilation block, before runtime cache allocation. - monkeypatch.setattr( - model_engine_module, - "get_attention_backend", - Mock(side_effect=RuntimeError("compile setup complete")), + engine._graphs = SimpleNamespace(no_cuda_graph=contextlib.nullcontext) + # The warmup resolves its chunk-alignment variant off the model's Mamba + # metadata class; a model declaring none takes the ``Mamba2Metadata`` + # default, i.e. no alignment rewrite. + engine.model = SimpleNamespace() + batch = object() + engine._dummy_requests = SimpleNamespace( + release_batch_context=lambda *_a, **_kw: _released_batch(batch) ) - engine = ModelEngine.__new__(ModelEngine) - with torch_compiling(False), pytest.raises(RuntimeError, match="compile setup complete"): - ModelEngine.__init__( - engine, - model_path="dummy", - mapping=mapping, - model=model, - llm_args=llm_args, - checkpoint_loader=Mock(), - ) - expected_prefill_only = ( - compile_enabled - and piecewise - and (compile_only_piecewise_graphs or model_kind in ("m3", "m3_vl")) - ) - assert engine._torch_compile_prefill_only is expected_prefill_only - if not compile_enabled: - compile_model.assert_not_called() - assert model.model is eager - else: - compile_model.assert_called_once_with( - eager, backend=engine._torch_compile_backend, fullgraph=True + cache_manager = _StandInMambaCacheManager() + resource_manager = SimpleNamespace( + get_resource_manager=lambda key: ( + cache_manager if key == ResourceManagerType.KV_CACHE_MANAGER else None ) - assert backend_factory.call_args.args[0] is False # Inductor stays disabled. - if expected_prefill_only: - assert isinstance(model.model, PrefillCompiledModel) - assert model.model.eager_model is eager - assert model.model.compiled_model is compiled - else: - assert model.model is compiled + ) + return engine, resource_manager -@pytest.mark.cpu_only -@pytest.mark.parametrize("local_contexts", [0, 1]) -def test_prefill_compile_uses_all_rank_prefill_decision( - monkeypatch: pytest.MonkeyPatch, - local_contexts: int, -) -> None: - """Route using the global prefill decision, regardless of local contexts.""" - eager = torch.nn.Linear(4, 4) - compiled = torch.nn.Module() - compiled.shared = eager - expected = torch.ones((2, 4)) - eager.forward = Mock(return_value=expected) - compiled.forward = Mock(return_value=expected) - router = PrefillCompiledModel(eager, compiled) - assert router.weight is eager.weight - assert list(router.parameters()) == list(eager.parameters()) - # Local decode-only ranks participate when another attention-DP rank - # prefills; an over-ceiling local context batch uses eager instead. - metadata = SimpleNamespace(num_contexts=local_contexts) - for eligible in (True, False): - monkeypatch.setattr( - decoder_runner_module, "get_per_request_prefill_cuda_graph_flag", lambda: eligible +def _run_mamba_warmup(engine: DecoderWarmup, resource_manager: object) -> None: + with ( + mock.patch.object(warmup_module, "MambaHybridCacheManager", _StandInMambaCacheManager), + mock.patch.object( + warmup_module.Mamba2Metadata, + "force_initial_states_for_warmup", + side_effect=contextlib.nullcontext, + ), + mock.patch.object(warmup_module, "clear_memory_buffers"), + ): + engine._run_mamba_hybrid_warmup(resource_manager) + + +class DecoderDsaWarmupTestCase(unittest.TestCase): + def test_warmup_builds_dsa_attn_metadata_when_no_forward_ran(self) -> None: + """The DSA top-K pre-compile hooks read attn_metadata; when every + warmup forward was skipped it is built on demand so the engines still + compile before serving.""" + from tensorrt_llm._torch.attention.backends.sparse.dsa import DSAtrtllmAttentionMetadata + + engine = object.__new__(DecoderWarmup) + engine._config = SimpleNamespace( + attention_backend=SimpleNamespace(Metadata=DSAtrtllmAttentionMetadata), + original_max_draft_len=2, + cuda_graph_batch_sizes=[8], + spec_config=None, ) - assert router(expected, attn_metadata=metadata) is expected - selected = compiled if eligible else eager - selected.forward.assert_called_once_with(expected, attn_metadata=metadata) - assert compiled.forward.call_count == eager.forward.call_count == 1 + engine._metadata = DecoderMetadata( + config=engine._config, + mapping=None, + model=None, + buffers=None, + attention_runtime_features=None, + ) + metadata = Mock(spec=DSAtrtllmAttentionMetadata) + def build( + kv_cache_manager: KVCacheManagerV2, draft_kv_cache_manager: KVCacheManagerV2 | None + ) -> Mock: + engine._metadata.attn_metadata = metadata + return metadata -@pytest.mark.cpu_only -def test_prefill_compile_preserves_partial_weight_reload( - monkeypatch: pytest.MonkeyPatch, -) -> None: - """Reload selected weights by original names into both model entry points.""" - import tensorrt_llm._torch.models.modeling_utils as modeling_utils - - eager = torch.nn.Sequential(OrderedDict(layers=torch.nn.Sequential(torch.nn.Linear(4, 4)))) - compiled = torch.compile(eager, backend="eager") - model = DecoderModelForCausalLM.__new__(DecoderModelForCausalLM) - torch.nn.Module.__init__(model) - model.model_config = SimpleNamespace( - pretrained_config=SimpleNamespace(tie_word_embeddings=False) - ) - model.model = PrefillCompiledModel(eager, compiled) - mapper = HfWeightMapper() - mapper._model = model - loader = ModelLoader.__new__(ModelLoader) - loader.weight_mapper = mapper - monkeypatch.setenv("TRT_LLM_DISABLE_LOAD_WEIGHTS_IN_PARALLEL", "True") - monkeypatch.setattr(modeling_utils, "local_mpi_rank", lambda: 0) - monkeypatch.setattr(torch.cuda, "set_device", lambda device: None) - monkeypatch.setattr(torch.cuda, "is_available", lambda: False) - - weight = eager.layers[0].weight - old_bias = eager.layers[0].bias.detach().clone() - replacement = torch.full_like(weight, 3) - for remove_duplicate in (False, True): - names = dict(model.named_modules(remove_duplicate=remove_duplicate)) - assert names["model.layers.0"] is eager.layers[0] - assert not any("eager_model" in name or "compiled_model" in name for name in names) - assert list(model.model.children()) == [eager] - assert list(model.model.state_dict()) == [ - "eager_model.layers.0.weight", - "eager_model.layers.0.bias", - ] - apply_tensor = Mock(side_effect=lambda tensor: tensor) - model.model._apply(apply_tensor) - assert apply_tensor.call_count == 2 - loader.reload(model, {"model.layers.0.weight": replacement}, allow_partial_loading=True) - - assert eager.layers[0].weight is weight - assert compiled.layers[0].weight is weight - torch.testing.assert_close(weight, replacement) - torch.testing.assert_close(eager.layers[0].bias, old_bias) - inputs = torch.ones(2, 4) - for eligible in (False, True): - monkeypatch.setattr( - decoder_runner_module, "get_per_request_prefill_cuda_graph_flag", lambda: eligible + engine._metadata.setup_attention = Mock(side_effect=build) + kv_cache_manager = Mock() + resource_manager = Mock() + resource_manager.get_resource_manager.return_value = kv_cache_manager + + engine._ensure_dsa_attn_metadata_for_warmup(resource_manager) + engine._warmup_cute_dsl_radix_topk() + + engine._metadata.setup_attention.assert_called_once_with(kv_cache_manager, None) + metadata.warmup_cute_dsl_radix_topk.assert_called_once_with(3) + metadata.warmup_selfsampling_topk.assert_called_once_with(3, batch_sizes=[8]) + + def test_warmup_does_not_build_metadata_for_non_dsa_backend(self) -> None: + engine = object.__new__(DecoderWarmup) + engine._config = SimpleNamespace( + attention_backend=SimpleNamespace(Metadata=TrtllmAttentionMetadata) + ) + engine._metadata = DecoderMetadata( + config=engine._config, + mapping=None, + model=None, + buffers=None, + attention_runtime_features=None, ) - torch.testing.assert_close(model.model(inputs), inputs @ replacement.t() + old_bias) + engine._metadata.setup_attention = Mock() + engine._ensure_dsa_attn_metadata_for_warmup(Mock()) + engine._warmup_cute_dsl_radix_topk() -@pytest.mark.cpu_only -@pytest.mark.parametrize("prefill_only", [False, True]) -@pytest.mark.parametrize("eligible", [False, True]) -@pytest.mark.parametrize("bypass", [False, True]) -@pytest.mark.parametrize("raises", [False, True]) -def test_prefill_compile_scopes_whole_model_forward( - monkeypatch: pytest.MonkeyPatch, - prefill_only: bool, - eligible: bool, - bypass: bool, - raises: bool, -) -> None: - """Scope compile dispatch by phase independently of the model wrapper.""" - observed = [] - - def forward(**kwargs: object) -> str: - """Record compile state as a stand-in for a model-specific epilogue.""" - observed.append(is_torch_compiling()) - # This represents work after the transformer, such as Eagle3 drafting. - if raises: - raise RuntimeError("epilogue failure") - observed.append(is_torch_compiling()) - return "done" - - eager_model = torch.nn.Identity() - eager_model.model_config = SimpleNamespace(extra_attrs={}) - eager_model.forward = forward - compiled_model = torch.nn.Identity() - compiled_model.forward = forward - prefill_compiled_model = PrefillCompiledModel(eager_model, compiled_model) - model = prefill_compiled_model if prefill_only else eager_model - engine = SimpleNamespace( - _model_caller=ModelCaller(model, prefill_compile_only=prefill_only), - _eager_workspace_reclaimer=None, - _moe_workspace_reclaimer=None, - ) - monkeypatch.setattr(model_call_module, "get_model_extra_attrs", lambda: {}) - monkeypatch.setattr( - model_call_module, "get_per_request_prefill_cuda_graph_flag", lambda: eligible - ) - monkeypatch.setattr(model_call_module, "is_trace_enabled", lambda name: False) - bypass_scope = ( - prefill_compiled_model.bypass() if prefill_only and bypass else contextlib.nullcontext() - ) - with torch_compiling(True), bypass_scope: - if raises: - with pytest.raises(RuntimeError, match="epilogue failure"): - DecoderRunner._model_forward(engine, is_dummy=False, attn_metadata=Mock()) - else: - assert ( - DecoderRunner._model_forward(engine, is_dummy=False, attn_metadata=Mock()) == "done" - ) - assert is_torch_compiling() - expected = eligible if prefill_only else True - assert observed == [expected] * (1 if raises else 2) + engine._metadata.setup_attention.assert_not_called() @pytest.mark.cpu_only @@ -423,9 +415,9 @@ def test_compiled_mxfp8_warmup_backend_selection( engine._model_caller = ModelCaller(engine.model, prefill_compile_only=prefill_only) cache = SimpleNamespace(get_num_available_tokens=lambda **kwargs: 16) resources = SimpleNamespace( - get_resource_manager=lambda key: cache - if key == ResourceManagerType.KV_CACHE_MANAGER - else None + get_resource_manager=lambda key: ( + cache if key == ResourceManagerType.KV_CACHE_MANAGER else None + ) ) tuner = Mock(profiling_cache={}) monkeypatch.setattr(decoder_warmup_module.AutoTuner, "get", lambda: tuner) @@ -530,9 +522,9 @@ def test_autotuner_warmup_skips_only_completed_token_shape( ) cache = SimpleNamespace(get_num_available_tokens=lambda **kwargs: 16) resources = SimpleNamespace( - get_resource_manager=lambda key: cache - if key == ResourceManagerType.KV_CACHE_MANAGER - else None + get_resource_manager=lambda key: ( + cache if key == ResourceManagerType.KV_CACHE_MANAGER else None + ) ) tuner = Mock(profiling_cache={}) get_tuner = Mock(return_value=tuner) @@ -561,297 +553,38 @@ def test_autotuner_warmup_skips_only_completed_token_shape( assert engine._completed_autotuner_warmup_num_tokens == 16 -@pytest.mark.parametrize("config_cls", [DraftTargetDecodingConfig, PARDDecodingConfig]) -def test_draft_length_resolution_resyncs_dynamic_buffers(config_cls): - config = config_cls(max_draft_len=3, speculative_model="dummy", draft_len_schedule={1: 3, 4: 2}) - engine = SimpleNamespace( - spec_config=config, - max_draft_len=3, - max_total_draft_tokens=config.tokens_per_gen_step - 1, - ) - requests = [ - SimpleNamespace( - py_draft_tokens=[7, 8, 9] + [0] * (config.tokens_per_gen_step - 4), - py_needs_onehot_draft_probs=False, - ) - for _ in range(2) - ] - batch = SimpleNamespace(batch_size=2, generation_requests=requests) - - def resolve(**kwargs): - return resolve_draft_len(config, batch, max_draft_len=3, static_draft_len=3, **kwargs) - - # Normal iteration selects K=2. - assert update_draft_len(engine, batch) == 2 - assert all(request.py_draft_tokens[:2] == [7, 8] for request in requests) - # Each explicit warmup shape resynchronizes the buffers and returns its - # length, even when the batch size is unchanged. - for draft_len in (0, 3, 1): - assert resolve(draft_len=draft_len) == draft_len - assert all( - len(request.py_draft_tokens) == config.get_runtime_tokens_per_gen_step(draft_len) - 1 - for request in requests - ) - assert resolve() == 2 - assert resolve(speculation_permanently_disabled=True) == 0 - assert all(request.py_draft_tokens == [] for request in requests) - - -def test_static_draft_length_resolution_preserves_proposals(): - request = SimpleNamespace(py_draft_tokens=[7], py_needs_onehot_draft_probs=False) - batch = SimpleNamespace(batch_size=1, generation_requests=[request]) - - assert resolve_draft_len(None, batch, max_draft_len=3, static_draft_len=5) == 5 - assert request.py_draft_tokens == [7] - assert not request.py_needs_onehot_draft_probs - - -# Minimal fixtures mirroring sibling test_pytorch_model_engine.py — duplicated -# rather than imported to keep this file self-contained and avoid sibling-test -# import fragility. -@dataclass -class _Config: - torch_dtype: torch.dtype - num_key_value_heads: int = 16 - num_attention_heads: int = 16 - hidden_size: int = 256 - architectures: list = None - - @property - def head_dim(self) -> int: - return self.hidden_size // self.num_attention_heads - - -class _DummyModel(torch.nn.Module): - def __init__(self, dtype: torch.dtype): - super().__init__() - self.model_config = ModelConfig(pretrained_config=_Config(torch_dtype=dtype)) - - def infer_max_seq_len(self): - return 2048 - - @property - def config(self): - return self.model_config.pretrained_config - - def forward(self, *args, **kwargs): - # Never actually called in these tests (the warmup helpers that would - # invoke forward are all patched), but must exist for engine init. - batch_size = kwargs["input_ids"].size(0) - return {"logits": torch.randn((batch_size, 10), device="cuda")} - - -class _DummyModelEngine(ModelEngine): - def __init__(self, llm_args: TorchLlmArgs, dtype: torch.dtype): - mapping = Mapping( - world_size=tensorrt_llm.mpi_world_size(), - tp_size=tensorrt_llm.mpi_world_size(), - rank=tensorrt_llm.mpi_rank(), - ) - super().__init__( - model_path="dummy", - mapping=mapping, - model=_DummyModel(dtype), - llm_args=llm_args, - ) - - -# KVCM2 requires a current CUDA context on whatever thread destroys it. A manager -# left for the cyclic GC is freed on whichever thread runs the collection -- under -# pytest's watchdog that thread has no context, and the first driver call in -# teardown fails with CUDA error 201, poisoning KVCM2 for the rest of the process. -# Releasing them here keeps teardown on the test's own thread. -_KV_CACHE_MANAGERS_TO_SHUTDOWN = [] - - -@pytest.fixture(autouse=True) -def _shutdown_kv_cache_managers(): - yield - while _KV_CACHE_MANAGERS_TO_SHUTDOWN: - _KV_CACHE_MANAGERS_TO_SHUTDOWN.pop().shutdown() - - -def _build_engine_and_resource_manager(): - tokens_per_block = 1 - max_tokens = 258 - num_layers = 1 - batch_size = 13 - llm_args = TorchLlmArgs( - model="dummy", - max_batch_size=batch_size, - max_num_tokens=max_tokens, - kv_cache_config=KvCacheConfig(max_tokens=max_tokens, use_kv_cache_manager_v2=True), - cuda_graph_config=CudaGraphConfig( - enable_padding=True, batch_sizes=[1, 2, 4, 8, 16, 32, 64, 128] - ), - ) - model_engine = _DummyModelEngine(llm_args, torch.half) - kv_cache_manager = KVCacheManagerV2( - llm_args.kv_cache_config, - tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF, - num_layers=num_layers, - num_kv_heads=model_engine.model.config.num_key_value_heads, - head_dim=model_engine.model.config.head_dim, - tokens_per_block=tokens_per_block, - max_seq_len=max_tokens, - max_batch_size=batch_size, - max_num_tokens=max_tokens, - mapping=Mapping(world_size=1, tp_size=1, rank=0), - dtype=tensorrt_llm.bindings.DataType.HALF, - ) - _KV_CACHE_MANAGERS_TO_SHUTDOWN.append(kv_cache_manager) - resource_manager = ResourceManager({ResourceManagerType.KV_CACHE_MANAGER: kv_cache_manager}) - return model_engine, resource_manager - - -@pytest.mark.parametrize("config_cls", [DraftTargetDecodingConfig, PARDDecodingConfig]) -@pytest.mark.parametrize( - "draft_len", [None, 0, 3, 1], ids=["general", "cuda_graph_0", "cuda_graph_3", "cuda_graph_1"] -) -def test_warmup_builders_resynchronize_stale_draft_buffers( - config_cls: type[DecodingBaseConfig], draft_len: int | None -) -> None: - config = config_cls(max_draft_len=3, speculative_model="dummy", draft_len_schedule={1: 3, 4: 2}) - engine, resource_manager = _build_engine_and_resource_manager() - runner = engine._runner - runner._config = replace( - runner._config, - spec_config=config, - max_draft_len=config.max_draft_len, - max_total_draft_tokens=config.tokens_per_gen_step - 1, - max_draft_loop_tokens=config.tokens_per_gen_step - 1, - ) - runner._get_runtime_tokens_per_gen_step = config.get_runtime_tokens_per_gen_step - requests = runner._dummy_requests - requests._config = runner._config - requests.get_runtime_tokens_per_gen_step = runner._get_runtime_tokens_per_gen_step - # Profiling a batch of two leaves K=2; every requested warmup shape below - # differs from that value and must synchronize the requests. - batch_size = 2 - - if draft_len is None: - expected_draft_len = runner._config.max_draft_len - warmup_request = requests.create_warmup_request( - resource_manager, - num_tokens=batch_size * config.tokens_per_gen_step, - num_gen_requests=batch_size, - ) - else: - expected_draft_len = draft_len - warmup_request = requests.create_cuda_graph_warmup_request( - resource_manager, batch_size, draft_len, max_seq_len=16 - ) - - with runner._dummy_requests.release_batch_context(warmup_request, resource_manager) as batch: - assert batch is not None - assert len(batch.generation_requests) == batch_size - expected_buffer_width = config.get_runtime_tokens_per_gen_step(expected_draft_len) - 1 - assert all( - len(request.py_draft_tokens) == expected_buffer_width - for request in batch.generation_requests - ) - - -@pytest.mark.parametrize("fails", [False, True], ids=["success", "failure"]) -def test_generation_capture_passes_local_speculation_state(fails: bool) -> None: - engine, resource_manager = _build_engine_and_resource_manager() - assert not engine.is_spec_decode - runner = engine._runner - batches = {} - - def create_request(resources, batch_size, draft_len, *args, **kwargs): - return batches.setdefault(draft_len, ScheduledRequests()) - - forward_warmup = Mock(side_effect=RuntimeError("capture failure") if fails else None) - with ( - patch.object(runner._graphs, "_get_graphs_to_capture", return_value=[(1, 0), (1, 2)]), - patch.object( - runner._dummy_requests, "create_cuda_graph_warmup_request", side_effect=create_request - ), - patch.object( - runner._dummy_requests, - "release_batch_context", - side_effect=lambda request, resources: contextlib.nullcontext(request), - ), - patch.object(runner._graphs, "_forward_warmup", forward_warmup), - ): - if fails: - with pytest.raises(RuntimeError, match="capture failure"): - runner._graphs._capture_generation_cuda_graphs(resource_manager) - else: - runner._graphs._capture_generation_cuda_graphs(resource_manager) - - expected_calls = [ - call(batches[2], resource_manager, enable_spec_decode=True, runtime_draft_len=2) - ] - if not fails: - expected_calls.append( - call(batches[0], resource_manager, enable_spec_decode=False, runtime_draft_len=0) - ) - assert forward_warmup.call_args_list == expected_calls - assert runner._graphs.force_lora_graph_for_capture is None - - @pytest.mark.parametrize("phase", ["general", "attention", "mamba", "prefill"]) def test_warmup_phases_pass_configured_speculation_state(phase: str) -> None: - engine, resource_manager = _build_engine_and_resource_manager() - runner = engine._runner - spec_config = DraftTargetDecodingConfig(max_draft_len=3, speculative_model="dummy") + spec = DraftTargetDecodingConfig(max_draft_len=3, speculative_model="dummy") + config = decoder_config( + spec_config=spec, + enable_autotuner=phase != "mamba", + prefill_cuda_graph_backend=PrefillCudaGraphBackend.BREAKABLE + if phase == "prefill" + else PrefillCudaGraphBackend.DISABLED, + prefill_cuda_graph_num_tokens=[4] if phase == "prefill" else [], + ) batch = ScheduledRequests() forward_warmup = Mock() - kv_cache_manager = resource_manager.get_resource_manager(ResourceManagerType.KV_CACHE_MANAGER) - - with contextlib.ExitStack() as stack: - spec_decode_config = replace(runner._config, spec_config=spec_config, is_spec_decode=True) - for owner in (runner._warmup, runner._graphs): - stack.enter_context(patch.object(owner, "_config", spec_decode_config)) - stack.enter_context( - patch.object(runner._dummy_requests, "create_warmup_request", Mock(return_value=batch)) - ) - stack.enter_context( - patch.object( - runner._dummy_requests, - "release_batch_context", - side_effect=lambda request, resources: contextlib.nullcontext(request), - ) - ) - for owner in (runner._warmup, runner._graphs): - stack.enter_context(patch.object(owner, "_forward_warmup", forward_warmup)) - if phase == "general": - runner._warmup._general_warmup_impl(resource_manager, [(4, 0)]) - elif phase == "attention": - runner._warmup._run_attention_warmup(resource_manager) - elif phase == "mamba": - stack.enter_context( - patch.object( - decoder_warmup_module, "MambaHybridCacheManager", type(kv_cache_manager) - ) - ) - stack.enter_context( - patch.object(kv_cache_manager, "get_num_available_tokens", return_value=8) - ) - stack.enter_context( - patch.object( - runner._warmup, - "_config", - replace(runner._warmup._config, enable_autotuner=False), - ) - ) - runner._warmup._run_mamba_hybrid_warmup(resource_manager) - else: - stack.enter_context( - patch.object( - runner._graphs, - "_config", - replace( - runner._graphs._config, - prefill_cuda_graph_backend=PrefillCudaGraphBackend.BREAKABLE, - prefill_cuda_graph_num_tokens=[4], - ), - ) - ) - runner._graphs._capture_prefill_cuda_graphs(resource_manager) - + requests = Mock() + requests.create_warmup_request.return_value = batch + requests.release_batch_context.side_effect = lambda request, resources: contextlib.nullcontext( + request + ) + warmup = decoder_warmup(config, dummy_requests=requests, forward_warmup=forward_warmup) + kv_cache_manager = Mock() + kv_cache_manager.get_num_available_tokens.return_value = 8 + resource_manager = Mock() + resource_manager.get_resource_manager.return_value = kv_cache_manager + if phase == "general": + warmup._general_warmup_impl(resource_manager, [(4, 0)]) + elif phase == "attention": + warmup._run_attention_warmup(resource_manager) + elif phase == "mamba": + with patch.object(decoder_warmup_module, "MambaHybridCacheManager", type(kv_cache_manager)): + warmup._run_mamba_hybrid_warmup(resource_manager) + else: + warmup._graphs._capture_prefill_cuda_graphs(resource_manager) jit_warmup = {"trtllm_gen_jit_warmup": True} if phase == "attention" else {} assert forward_warmup.call_count > 0 assert ( @@ -865,432 +598,145 @@ def test_warmup_phases_pass_configured_speculation_state(phase: str) -> None: ) -class _Tracker: - """Records method-call order (and each call's arguments) via mock side_effects.""" - - def __init__(self): - self.calls = [] - self.call_args = {} +@pytest.mark.cpu_only +@pytest.mark.parametrize("tp_size", [1, 2]) +def test_indexer_q_prewarm_runs_every_fallback_tactic(tp_size: int) -> None: + indexer = _deepseek_v4_indexer(q_lora_rank=256) + owner = _indexer_q_prewarm_owner([torch.nn.Linear(2, 2), indexer], tp_size) + events = [] + indexer._project_and_quantize_q.side_effect = lambda *args: events.append("compile") + owner.dist.tp_allgather.side_effect = lambda *args: events.append("barrier") - def __call__(self, name): - def _wrapped(*args, **kwargs): - self.calls.append(name) - self.call_args.setdefault(name, []).append((args, kwargs)) + with patch.object(torch.cuda, "synchronize", side_effect=lambda: events.append("synchronize")): + DecoderWarmup._prewarm_cute_dsl_indexer_q(owner) - return _wrapped + calls = indexer._project_and_quantize_q.call_args_list + assert [c.args[0].shape for c in calls] == [(4, 256), (8, 256), (16, 256)] + assert [c.args[1].shape for c in calls] == [(4,), (8,), (16,)] + # The fused path only accepts bf16 activations. + assert all(c.args[0].dtype == torch.bfloat16 for c in calls) + assert all(c.args[1].dtype == torch.int32 for c in calls) + indexer._is_fused_project_mxfp4_enabled.assert_called_once_with(torch.bfloat16) + # The barrier is only useful once every compile on this rank has finished. + expected_barrier = ["barrier"] if tp_size > 1 else [] + assert events == ["compile"] * 3 + ["synchronize"] + expected_barrier + if tp_size > 1: + owner.dist.tp_allgather.assert_called_once_with(1) -def _run_warmup_tracked( - model_engine, resource_manager, *, force_helix_cp=False, capture_logs=False, tracker=None -): - """Patch warmup stages and cleanup operations, then run warmup. +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "case", ["no_indexer", "indexer_without_wq_b", "indexer_of_other_pp_stage"] +) +def test_indexer_q_prewarm_is_noop_without_local_dsv4_indexer(case: str) -> None: + if case == "indexer_of_other_pp_stage": + indexer = _other_pp_stage_indexer() + else: + indexer = _deepseek_v4_indexer(with_wq_b=False) + model_modules = [torch.nn.Linear(2, 2)] + if case != "no_indexer": + model_modules.append(indexer) + owner = _indexer_q_prewarm_owner(model_modules, tp_size=2) - Optionally force helix CP and capture logs. Pass ``tracker`` to inspect the - recorded call arguments. Returns (call_order_list, log_records_or_None). - """ - if tracker is None: - tracker = _Tracker() - # Engines left by earlier tests are cyclic; finalize them before tracking. - gc.collect() - # Observe this module's explicit cleanup without counting unrelated - # destructors triggered by gc.collect() or lazy imports. - tracked_cuda = SimpleNamespace(**vars(torch.cuda)) - tracked_cuda.empty_cache = tracker("empty_cache") - tracked_torch = SimpleNamespace(**vars(torch)) - tracked_torch.cuda = tracked_cuda - helix_ctx = ( - patch.object(model_engine.mapping, "has_cp_helix", return_value=True) - if force_helix_cp - else contextlib.nullcontext() - ) + with patch.object(torch.cuda, "synchronize") as synchronize: + DecoderWarmup._prewarm_cute_dsl_indexer_q(owner) - with ( - helix_ctx, - patch.object( - model_engine._runner._warmup, "_general_warmup", side_effect=tracker("general_warmup") - ), - patch.object( - model_engine._runner._warmup, "_run_autotuner_warmup", side_effect=tracker("autotuner") - ), - patch.object( - model_engine._runner._graphs, - "_capture_generation_cuda_graphs", - side_effect=tracker("generation_cuda_graph"), - ), - patch.object( - model_engine._runner._graphs, - "_capture_additional_cuda_graphs", - side_effect=tracker("additional_cuda_graph"), - ), - patch.object( - model_engine._runner._graphs, - "_capture_prefill_cuda_graphs", - side_effect=tracker("prefill_cuda_graph"), - ), - patch( - "tensorrt_llm._torch.pyexecutor.engine.runners.decoder.warmup.warmup_sampling_module", - side_effect=tracker("sampling_warmup"), - ), - patch.object(decoder_warmup_module, "torch", tracked_torch), - patch( - "tensorrt_llm._torch.custom_ops.torch_custom_ops.MoERunner.clear_all_workspaces", - side_effect=tracker("moe_clear"), - ), - ): - if capture_logs: - with _capture_tllm_logs() as logs: - model_engine.warmup(resource_manager) - return tracker.calls, logs - model_engine.warmup(resource_manager) - return tracker.calls, None + indexer._project_and_quantize_q.assert_not_called() + synchronize.assert_not_called() + owner.dist.tp_allgather.assert_not_called() -@contextlib.contextmanager -def _capture_tllm_logs(): - """Capture logger.info calls emitted from the decoder runner module. - - tensorrt_llm.logger.logger is a custom Singleton (not a stdlib - logging.Logger) and does not route through stdlib logging by default, - so a logging.Handler attached to logging.getLogger("tensorrt_llm") - sees nothing. Patch the logger.info bound on the decoder runner module - directly so we observe exactly the messages warmup() emits. - """ - records = [] +@pytest.mark.cpu_only +def test_indexer_q_prewarm_without_fused_path_compiles_nothing_but_reaches_barrier() -> None: + """The fused-path check is rank-local, so it must not decide the TP barrier.""" + indexer = _deepseek_v4_indexer(fused_enabled=False) + owner = _indexer_q_prewarm_owner([indexer], tp_size=2) - def _record(*msg): - records.append(" ".join(str(m) for m in msg)) + with ( + patch.object(torch.cuda, "synchronize") as synchronize, + patch.object(decoder_warmup_module.logger, "info") as info, + ): + DecoderWarmup._prewarm_cute_dsl_indexer_q(owner) - with patch.object(decoder_warmup_module.logger, "info", side_effect=_record): - yield records + indexer._project_and_quantize_q.assert_not_called() + synchronize.assert_not_called() + info.assert_not_called() + owner.dist.tp_allgather.assert_called_once_with(1) @pytest.mark.cpu_only -@pytest.mark.parametrize("schedule", [{4: 4, 8: 2, 16: 1}, {4: 2}, {16: 4}, {4: 0}]) -def test_dynamic_warmup_graphs_match_replay_schedule(schedule: dict[int, int]) -> None: - from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import CUDAGraphRunner - from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.graphs import ( - compute_dynamic_draft_len_mapping, - ) - from tensorrt_llm._torch.speculative.utils import get_draft_len_for_batch_size - - batch_sizes = [1, 2, 4, 8, 16] - runner = object.__new__(DecoderCudaGraphs) - mode = SimpleNamespace(support_dynamic_draft_len=lambda: True) - runner._config = SimpleNamespace( - spec_config=SimpleNamespace(draft_len_schedule=schedule, spec_dec_mode=mode), - cuda_graph_batch_sizes=batch_sizes, - original_max_draft_len=max(schedule.values()), +def test_indexer_q_prewarm_compiles_first_fused_indexer() -> None: + unfused_indexer = _deepseek_v4_indexer(fused_enabled=False) + fused_indexer = _deepseek_v4_indexer() + later_fused_indexer = _deepseek_v4_indexer() + owner = _indexer_q_prewarm_owner( + [unfused_indexer, fused_indexer, later_fused_indexer], tp_size=2 ) - runner._dynamic_draft_len_mapping = compute_dynamic_draft_len_mapping(runner._config) - graphs = runner._get_graphs_to_capture(batch_sizes) - replay = object.__new__(CUDAGraphRunner) - replay.supported_batch_sizes = batch_sizes - replay.dynamic_draft_len_mapping = runner._dynamic_draft_len_mapping - - reachable_graphs = set() - for batch_size in range(1, max(batch_sizes) + 1): - draft_len = get_draft_len_for_batch_size(schedule, batch_size, max(schedule.values())) - padded_size = replay._round_up_batch_size_with_draft_len(batch_size, draft_len) - assert padded_size >= batch_size - reachable_graphs.add((padded_size, draft_len)) - - assert set(graphs) == reachable_graphs - assert len(graphs) == len(reachable_graphs) - - -class TestWarmupCleanup(unittest.TestCase): - """Lock in warmup-cleanup behavior introduced by PR #14609 (Plan B).""" - - @pytest.mark.cpu_only - def test_no_kv_cache_warmup_delegates_runner_lifecycle(self): - """Verify no-KV-cache warmup delegates warmup and graph capture to the runner.""" - model_engine = object.__new__(ModelEngine) - model_engine.model = SimpleNamespace(model_config=SimpleNamespace(is_encoder_decoder=False)) - model_engine.moe_load_balancer = None - model_engine._metrics = {} - model_engine._warmup_depth = 0 - model_engine._runner = Mock(spec=NoKVCacheRunner) - resource_manager = Mock() - resource_manager.get_resource_manager.return_value = None - - model_engine.warmup(resource_manager) - - self.assertEqual( - model_engine._runner.method_calls, - [call.warmup(resource_manager)], - ) - - @pytest.mark.cpu_only - def test_no_kv_cache_warmup_rejects_allocated_kv_cache(self): - """Verify a no-KV-cache runner rejects an allocated KV cache before warmup.""" - model_engine = object.__new__(ModelEngine) - model_engine._metrics = {} - model_engine.model = SimpleNamespace(model_config=SimpleNamespace(is_encoder_decoder=False)) - model_engine.moe_load_balancer = None - model_engine._warmup_depth = 0 - model_engine._runner = Mock(spec=NoKVCacheRunner) - runner = model_engine._runner - runner._validate_resources = NoKVCacheRunner._validate_resources.__get__(runner) - runner.warmup.side_effect = NoKVCacheRunner.warmup.__get__(runner) - resource_manager = Mock() - resource_manager.get_resource_manager.return_value = object() - - with self.assertRaisesRegex( - AssertionError, - "no-KV-cache runner was initialized, but a KV cache manager was allocated", - ): - model_engine.warmup(resource_manager) - model_engine._runner.warmup.assert_called_once_with(resource_manager) + with patch.object(torch.cuda, "synchronize"): + DecoderWarmup._prewarm_cute_dsl_indexer_q(owner) - @pytest.mark.cpu_only - def test_decoder_runner_warmup_sampling_and_kv_cache_cleanup(self) -> None: - """Verify sampling warmup coverage and KV-cache cleanup timing within total warmup.""" - runner = object.__new__(DecoderRunner) - warmup = runner._warmup = object.__new__(DecoderWarmup) - warmup._warmup_timer = _WarmupTimer(rank=0) - warmup.model = SimpleNamespace(config=SimpleNamespace(vocab_size=128)) - warmup._config = SimpleNamespace( - dtype=torch.float16, cuda_graph_batch_sizes=[1, 4], enable_in_graph_sampling=False - ) - model_engine = object.__new__(ModelEngine) - model_engine._metrics = warmup._metrics = {} - model_engine.moe_load_balancer = None - model_engine._warmup_depth = 0 - model_engine._runner = runner - resource_manager = Mock() - resource_manager.get_resource_manager.return_value = None - events = [] - phase = warmup._warmup_timer.phase + unfused_indexer._project_and_quantize_q.assert_not_called() + assert fused_indexer._project_and_quantize_q.call_count == 3 + later_fused_indexer._project_and_quantize_q.assert_not_called() + owner.dist.tp_allgather.assert_called_once_with(1) - @contextlib.contextmanager - def record_metric_scope( - name: str, *, metrics: dict[str, float], metric_name: str - ) -> Iterator[None]: - events.append(("enter", metric_name)) - with phase(name, metrics=metrics, metric_name=metric_name): - yield - events.append(("exit", metric_name)) - with ( - patch.object( - warmup._warmup_timer, - "phase", - side_effect=record_metric_scope, - ), - patch( - "tensorrt_llm._torch.pyexecutor.engine.runners.decoder.warmup.warmup_sampling_module", - side_effect=lambda: events.append("sampling"), - ) as warmup_sampling, - _capture_tllm_logs() as logs, - ): - for enabled in (False, True): - with self.subTest(enable_in_graph_sampling=enabled): - events.clear() - warmup_sampling.reset_mock() - warmup._config.enable_in_graph_sampling = enabled - runner._eager_workspace_reclaimer = object() - model_engine._warmup_impl(resource_manager) - self.assertEqual( - events, - [ - ("enter", "sampling_warmup_seconds"), - "sampling", - ("exit", "sampling_warmup_seconds"), - ], - ) - self.assertIsNone(runner._eager_workspace_reclaimer) - warmup_sampling.assert_called_once_with() +@pytest.mark.cpu_only +def test_indexer_q_prewarm_failure_still_reaches_barrier() -> None: + """A per-shape failure must not skip the TP barrier the other ranks wait on.""" + indexer = _deepseek_v4_indexer() + indexer._project_and_quantize_q.side_effect = [RuntimeError("compile"), None, None] + owner = _indexer_q_prewarm_owner([indexer], tp_size=2) - self.assertTrue( - any("Skipping warm up as no KV Cache manager allocated." in log for log in logs) - ) - self.assertGreaterEqual(model_engine.metrics["sampling_warmup_seconds"], 0) + with ( + patch.object(torch.cuda, "synchronize"), + patch.object(decoder_warmup_module.logger, "warning") as warning, + ): + DecoderWarmup._prewarm_cute_dsl_indexer_q(owner) - kv_cache_manager = Mock() - kv_cache_manager.check_invalid_values_in_kv_cache.return_value = False - resource_manager.get_resource_manager.return_value = kv_cache_manager - model_engine._metrics.clear() - with ( - patch.object(model_engine, "_warmup_impl") as warmup_impl, - patch("time.perf_counter", side_effect=[0.0, 2.0, 5.0, 7.0]), - ): - model_engine.warmup(resource_manager) + assert indexer._project_and_quantize_q.call_count == 3 + warning.assert_called_once() + owner.dist.tp_allgather.assert_called_once_with(1) - warmup_impl.assert_called_once_with(resource_manager) - kv_cache_manager.check_invalid_values_in_kv_cache.assert_called_once_with( - fill_with_zero=True - ) - self.assertEqual( - model_engine.metrics, - {"kv_cache_cleanup_seconds": 3.0, "total_warmup_seconds": 7.0}, - ) - def test_encoder_decoder_encoder_warmup_delegates_runner_lifecycle(self): - model_engine = object.__new__(ModelEngine) - model_engine.model = SimpleNamespace(model_config=SimpleNamespace(is_encoder_decoder=True)) - model_engine.moe_load_balancer = None - model_engine._warmup_depth = 0 - model_engine._runner = Mock(spec=EncoderDecoderRunner) - warmup_depths = [] - model_engine._runner.warmup_encoder.side_effect = lambda _: warmup_depths.append( - model_engine._warmup_depth - ) - resource_manager = object() +@pytest.mark.cpu_only +def test_indexer_q_prewarm_does_not_report_completion_when_every_compile_fails() -> None: + indexer = _deepseek_v4_indexer() + indexer._project_and_quantize_q.side_effect = RuntimeError("compile") + owner = _indexer_q_prewarm_owner([indexer], tp_size=2) - model_engine.warmup_encoder(resource_manager) + with ( + patch.object(torch.cuda, "synchronize"), + patch.object(decoder_warmup_module.logger, "warning") as warning, + patch.object(decoder_warmup_module.logger, "info") as info, + ): + DecoderWarmup._prewarm_cute_dsl_indexer_q(owner) - self.assertEqual( - model_engine._runner.method_calls, - [call.warmup_encoder(resource_manager)], - ) - self.assertEqual(warmup_depths, [1]) - self.assertEqual(model_engine._warmup_depth, 0) + assert warning.call_count == 3 + info.assert_not_called() + owner.dist.tp_allgather.assert_called_once_with(1) - @pytest.mark.cpu_only - def test_cuda_graph_metrics_exclude_piecewise_stages(self) -> None: - """Verify generation graph timing excludes prefill and includes LoRA warmup cleanup.""" - graphs = DecoderCudaGraphs( - cuda_graph_runner=SimpleNamespace(enabled=True, is_warmup_only=False), - breakable_cuda_graph_runner=None, - config=SimpleNamespace( - enable_autotuner=True, prefill_cuda_graph_backend=PrefillCudaGraphBackend.PIECEWISE - ), - mapping=Mapping(), - model=SimpleNamespace(modules=lambda: []), - sparse_attention_config=None, - metrics={}, - dummy_requests=None, - agreement=None, - metadata=None, - get_runtime_tokens_per_gen_step=None, - dynamic_draft_len_mapping=None, - lora=SimpleNamespace(cuda_graph_manager=object()), - forward_warmup=None, - torch_compile_bypass=contextlib.nullcontext, - context_promotion_supported=True, - ) - resource_manager = ResourceManager({}) - events = [] - @contextlib.contextmanager - def record_metric_scope(metric_name: str, _metrics: dict[str, float]) -> Iterator[None]: - events.append(("enter", metric_name)) - try: - yield - finally: - events.append(("exit", metric_name)) +@pytest.mark.cpu_only +def test_indexer_q_prewarm_skips_indexers_of_other_pp_stages() -> None: + remote_indexer = _other_pp_stage_indexer() + local_indexer = _deepseek_v4_indexer() + owner = _indexer_q_prewarm_owner([remote_indexer, local_indexer], tp_size=2) - autotuner = SimpleNamespace( - cache_pp_recv=lambda: events.append(("lora", "cache_pp_recv")), - cache_pp_send=lambda: events.append(("lora", "cache_pp_send")), - clean_pp_flag=lambda: events.append(("lora", "clean_pp_flag")), - ) - lora_cleanup_events = [ - ("lora", "cache_pp_recv"), - ("lora", "cache_pp_send"), - ("lora", "clean_pp_flag"), - ("exit", "lora_autotune"), - ("exit", "gen_cuda_graph_warmup_seconds"), - ] + with patch.object(torch.cuda, "synchronize"): + DecoderWarmup._prewarm_cute_dsl_indexer_q(owner) - with ( - patch( - "tensorrt_llm._torch.pyexecutor.engine.runners.decoder.graphs.autotune", - side_effect=lambda **_: record_metric_scope("lora_autotune", {}), - ), - patch( - "tensorrt_llm._torch.pyexecutor.engine.runners.decoder.graphs.AutoTuner.get", - return_value=autotuner, - ), - patch( - "tensorrt_llm._torch.modules.linear.flashinfer_mxfp8_decode_graph_capture", - side_effect=contextlib.nullcontext, - ), - patch( - "tensorrt_llm._torch.pyexecutor.engine.runners.decoder.graphs.timing_metric", - side_effect=record_metric_scope, - ), - patch.object( - graphs, - "_capture_generation_cuda_graphs", - side_effect=lambda _: events.append(("stage", "generation")), - ) as generation, - patch.object( - graphs, - "_capture_additional_cuda_graphs", - side_effect=lambda _: events.append(("stage", "additional")), - ), - patch.object( - graphs, - "_capture_prefill_cuda_graphs", - side_effect=lambda _: events.append(("stage", "piecewise")), - ) as piecewise, - ): - graphs.run_cuda_graph_warmup(resource_manager) - - self.assertEqual( - events, - [ - ("enter", "gen_cuda_graph_capture_seconds"), - ("stage", "generation"), - ("stage", "additional"), - ("exit", "gen_cuda_graph_capture_seconds"), - ("stage", "piecewise"), - ], - ) + remote_indexer._project_and_quantize_q.assert_not_called() + assert local_indexer._project_and_quantize_q.call_count == 3 + owner.dist.tp_allgather.assert_called_once_with(1) - events.clear() - piecewise.reset_mock() - graphs.cuda_graph_runner.is_warmup_only = True - graphs.run_cuda_graph_warmup(resource_manager) - - self.assertEqual( - events, - [ - ("enter", "gen_cuda_graph_warmup_seconds"), - ("enter", "lora_autotune"), - ("stage", "generation"), - ("stage", "additional"), - *lora_cleanup_events, - ], - ) - piecewise.assert_not_called() - - events.clear() - generation.side_effect = RuntimeError("warmup failed") - with self.assertRaisesRegex(RuntimeError, "warmup failed"): - graphs.run_cuda_graph_warmup(resource_manager) - self.assertEqual( - events, - [ - ("enter", "gen_cuda_graph_warmup_seconds"), - ("enter", "lora_autotune"), - *lora_cleanup_events, - ], - ) - piecewise.assert_not_called() - - events.clear() - generation.reset_mock() - graphs.cuda_graph_runner.enabled = False - graphs._config.prefill_cuda_graph_backend = PrefillCudaGraphBackend.DISABLED - graphs.run_cuda_graph_warmup(resource_manager) - self.assertEqual( - events, - [ - ("enter", "gen_cuda_graph_warmup_seconds"), - ("enter", "lora_autotune"), - *lora_cleanup_events, - ], - ) - generation.assert_not_called() - piecewise.assert_not_called() - def test_empty_cache_fires_immediately_after_autotuner(self): +class TestWarmupCleanup(unittest.TestCase): + def test_empty_cache_fires_immediately_after_autotuner(self) -> None: """Change 1 placement: empty_cache must be the call right after _run_autotuner_warmup.""" - model_engine, resource_manager = _build_engine_and_resource_manager() + model_engine, resource_manager = self.enterContext(_warmup_and_resources()) calls, _ = _run_warmup_tracked(model_engine, resource_manager) self.assertIn("autotuner", calls) @@ -1306,11 +752,11 @@ def test_empty_cache_fires_immediately_after_autotuner(self): f"Expected empty_cache right after autotuner; got {calls}", ) - def test_empty_cache_count_under_default(self): + def test_empty_cache_count_under_default(self) -> None: """Default warmup should call empty_cache exactly twice: once at the end of step (a) (pre-existing) and once after step (b) (Change 1).""" - model_engine, resource_manager = _build_engine_and_resource_manager() + model_engine, resource_manager = self.enterContext(_warmup_and_resources()) calls, _ = _run_warmup_tracked(model_engine, resource_manager) self.assertEqual( calls.count("empty_cache"), @@ -1318,11 +764,11 @@ def test_empty_cache_count_under_default(self): f"Expected exactly 2 empty_cache calls; got order={calls}", ) - def test_helix_cp_runs_only_the_generation_autotuner_warmup(self): + def test_helix_cp_runs_only_the_generation_autotuner_warmup(self) -> None: """With Helix CP, can_run_general_warmup is False and the context warmups (general, Mamba hybrid) are skipped; the autotuner warmup still runs, on generation requests, followed by one empty_cache.""" - model_engine, resource_manager = _build_engine_and_resource_manager() + model_engine, resource_manager = self.enterContext(_warmup_and_resources()) tracker = _Tracker() calls, _ = _run_warmup_tracked( model_engine, resource_manager, force_helix_cp=True, tracker=tracker @@ -1342,7 +788,7 @@ def test_helix_cp_runs_only_the_generation_autotuner_warmup(self): self.assertEqual(calls[calls.index("autotuner") + 1], "empty_cache") @pytest.mark.cpu_only - def test_generation_only_autotuner_warmup_tunes_one_full_generation_batch(self): + def test_generation_only_autotuner_warmup_tunes_one_full_generation_batch(self) -> None: """Helix CP tunes on max_batch_size generation requests (bounded by max_num_tokens), without the context batch.""" tuner = SimpleNamespace( @@ -1387,9 +833,9 @@ def test_generation_only_autotuner_warmup_tunes_one_full_generation_batch(self): ) kv_cache_manager = SimpleNamespace(get_num_available_tokens=get_num_available_tokens) resource_manager = SimpleNamespace( - get_resource_manager=lambda key: kv_cache_manager - if key == ResourceManagerType.KV_CACHE_MANAGER - else None + get_resource_manager=lambda key: ( + kv_cache_manager if key == ResourceManagerType.KV_CACHE_MANAGER else None + ) ) with ( @@ -1410,12 +856,12 @@ def test_generation_only_autotuner_warmup_tunes_one_full_generation_batch(self): self.assertEqual(engine._completed_autotuner_warmup_num_tokens, 100) @pytest.mark.cpu_only - def test_flashinfer_mxfp8_respects_disabled_global_autotuner(self): + def test_flashinfer_mxfp8_respects_disabled_global_autotuner(self) -> None: """The global autotuner switch also disables automatic FlashInfer tuning.""" calls = [] @contextlib.contextmanager - def flashinfer_autotune(): + def flashinfer_autotune() -> Iterator[None]: calls.append("flashinfer_autotune_enter") yield calls.append("flashinfer_autotune_exit") @@ -1462,19 +908,19 @@ def flashinfer_autotune(): flashinfer_module.autotune.assert_not_called() @pytest.mark.cpu_only - def test_mxfp8_native_and_flashinfer_use_separate_warmup_passes(self): + def test_mxfp8_native_and_flashinfer_use_separate_warmup_passes(self) -> None: """Native and FlashInfer backends each receive an isolated tuning forward.""" calls = [] @contextlib.contextmanager - def trtllm_autotune(**kwargs): + def trtllm_autotune(**kwargs: object) -> Iterator[None]: self.assertIsNone(kwargs["cache_path"]) calls.append("trtllm_autotune_enter") yield calls.append("trtllm_autotune_exit") @contextlib.contextmanager - def flashinfer_autotune(): + def flashinfer_autotune() -> Iterator[None]: calls.append("flashinfer_autotune_enter") yield calls.append("flashinfer_autotune_exit") @@ -1552,9 +998,9 @@ def flashinfer_autotune(): ) kv_cache_manager = SimpleNamespace(get_num_available_tokens=lambda **kwargs: 16) resource_manager = SimpleNamespace( - get_resource_manager=lambda key: kv_cache_manager - if key == ResourceManagerType.KV_CACHE_MANAGER - else None + get_resource_manager=lambda key: ( + kv_cache_manager if key == ResourceManagerType.KV_CACHE_MANAGER else None + ) ) with ( @@ -1597,19 +1043,19 @@ def flashinfer_autotune(): tuner.setup_distributed_state.assert_called_with(engine.mapping, engine.dist) @pytest.mark.cpu_only - def test_native_mxfp8_falls_back_after_missing_warmup_batch(self): + def test_native_mxfp8_falls_back_after_missing_warmup_batch(self) -> None: """A missing startup batch latches native MXFP8 to the default tactic.""" calls = [] @contextlib.contextmanager - def trtllm_autotune(**kwargs): + def trtllm_autotune(**kwargs: object) -> Iterator[None]: self.assertIsNone(kwargs["cache_path"]) calls.append("autotune_enter") yield calls.append("autotune_exit") @contextlib.contextmanager - def flashinfer_autotune(): + def flashinfer_autotune() -> Iterator[None]: calls.append("flashinfer_autotune_enter") yield calls.append("flashinfer_autotune_exit") @@ -1682,9 +1128,9 @@ def flashinfer_autotune(): ) kv_cache_manager = SimpleNamespace(get_num_available_tokens=lambda **kwargs: 16) resource_manager = SimpleNamespace( - get_resource_manager=lambda key: kv_cache_manager - if key == ResourceManagerType.KV_CACHE_MANAGER - else None + get_resource_manager=lambda key: ( + kv_cache_manager if key == ResourceManagerType.KV_CACHE_MANAGER else None + ) ) with ( @@ -1722,7 +1168,7 @@ def flashinfer_autotune(): engine._forward_warmup.assert_not_called() @pytest.mark.cpu_only - def test_flashinfer_mxfp8_rank_mismatch_falls_back_before_warmup(self): + def test_flashinfer_mxfp8_rank_mismatch_falls_back_before_warmup(self) -> None: """TP and PP ranks agree on fallback before the tuning forward.""" flashinfer_module = ModuleType("flashinfer") flashinfer_module.mm_mxfp8 = Mock() @@ -1790,9 +1236,9 @@ def test_flashinfer_mxfp8_rank_mismatch_falls_back_before_warmup(self): ) kv_cache_manager = SimpleNamespace(get_num_available_tokens=lambda **kwargs: 16) resource_manager = SimpleNamespace( - get_resource_manager=lambda key: kv_cache_manager - if key == ResourceManagerType.KV_CACHE_MANAGER - else None + get_resource_manager=lambda key: ( + kv_cache_manager if key == ResourceManagerType.KV_CACHE_MANAGER else None + ) ) with ( @@ -1823,7 +1269,7 @@ def test_flashinfer_mxfp8_rank_mismatch_falls_back_before_warmup(self): self.assertEqual(engine._forward_warmup.call_count, 1) @pytest.mark.cpu_only - def test_native_mxfp8_respects_disabled_global_autotuner(self): + def test_native_mxfp8_respects_disabled_global_autotuner(self) -> None: """Avoid native MXFP8 warmup when the global autotuner is disabled.""" with ( patch( @@ -1847,170 +1293,156 @@ def test_native_mxfp8_respects_disabled_global_autotuner(self): self.assertFalse(method.needs_native_autotune) -def _indexer_q_prewarm_owner(modules: list, tp_size: int) -> SimpleNamespace: - return SimpleNamespace( - model=SimpleNamespace(modules=lambda: iter(modules)), - mapping=SimpleNamespace(tp_size=tp_size), - dist=Mock(), - ) - - -def _deepseek_v4_indexer( - q_lora_rank: int = 256, with_wq_b: bool = True, fused_enabled: bool = True -) -> Mock: - from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4 import DeepseekV4Indexer - - indexer = Mock(spec=DeepseekV4Indexer) - indexer.wq_b = ( - SimpleNamespace(weight=torch.empty((64, q_lora_rank), dtype=torch.bfloat16)) - if with_wq_b - else None - ) - indexer._is_fused_project_mxfp4_enabled.return_value = fused_enabled - return indexer - - -def _other_pp_stage_indexer() -> Mock: - """Layers owned by another pipeline stage keep wq_b but lose its weight.""" - indexer = _deepseek_v4_indexer() - indexer.wq_b = torch.nn.Linear(256, 64, bias=False) - remove_weights(indexer.wq_b) - return indexer - - -@pytest.mark.cpu_only -@pytest.mark.parametrize("tp_size", [1, 2]) -def test_indexer_q_prewarm_runs_every_fallback_tactic(tp_size: int) -> None: - indexer = _deepseek_v4_indexer(q_lora_rank=256) - owner = _indexer_q_prewarm_owner([torch.nn.Linear(2, 2), indexer], tp_size) - events = [] - indexer._project_and_quantize_q.side_effect = lambda *args: events.append("compile") - owner.dist.tp_allgather.side_effect = lambda *args: events.append("barrier") - - with patch.object(torch.cuda, "synchronize", side_effect=lambda: events.append("synchronize")): - DecoderWarmup._prewarm_cute_dsl_indexer_q(owner) - - calls = indexer._project_and_quantize_q.call_args_list - assert [c.args[0].shape for c in calls] == [(4, 256), (8, 256), (16, 256)] - assert [c.args[1].shape for c in calls] == [(4,), (8,), (16,)] - # The fused path only accepts bf16 activations. - assert all(c.args[0].dtype == torch.bfloat16 for c in calls) - assert all(c.args[1].dtype == torch.int32 for c in calls) - indexer._is_fused_project_mxfp4_enabled.assert_called_once_with(torch.bfloat16) - # The barrier is only useful once every compile on this rank has finished. - expected_barrier = ["barrier"] if tp_size > 1 else [] - assert events == ["compile"] * 3 + ["synchronize"] + expected_barrier - if tp_size > 1: - owner.dist.tp_allgather.assert_called_once_with(1) - - @pytest.mark.cpu_only @pytest.mark.parametrize( - "case", ["no_indexer", "indexer_without_wq_b", "indexer_of_other_pp_stage"] + "world_size,dwdp_size,is_fatal", + [(1, 0, False), (2, 0, True), (1, 4, True)], ) -def test_indexer_q_prewarm_is_noop_without_local_dsv4_indexer(case: str) -> None: - if case == "indexer_of_other_pp_stage": - indexer = _other_pp_stage_indexer() - else: - indexer = _deepseek_v4_indexer(with_wq_b=False) - model_modules = [torch.nn.Linear(2, 2)] - if case != "no_indexer": - model_modules.append(indexer) - owner = _indexer_q_prewarm_owner(model_modules, tp_size=2) - - with patch.object(torch.cuda, "synchronize") as synchronize: - DecoderWarmup._prewarm_cute_dsl_indexer_q(owner) - - indexer._project_and_quantize_q.assert_not_called() - synchronize.assert_not_called() - owner.dist.tp_allgather.assert_not_called() - - -@pytest.mark.cpu_only -def test_indexer_q_prewarm_without_fused_path_compiles_nothing_but_reaches_barrier() -> None: - """The fused-path check is rank-local, so it must not decide the TP barrier.""" - indexer = _deepseek_v4_indexer(fused_enabled=False) - owner = _indexer_q_prewarm_owner([indexer], tp_size=2) - - with ( - patch.object(torch.cuda, "synchronize") as synchronize, - patch.object(decoder_warmup_module.logger, "info") as info, - ): - DecoderWarmup._prewarm_cute_dsl_indexer_q(owner) +def test_general_warmup_oom_policy(world_size: int, dwdp_size: int, is_fatal: bool) -> None: + engine = _general_warmup_engine(world_size=world_size, dwdp_size=dwdp_size) + error = torch.OutOfMemoryError("asymmetric OOM") + engine._forward_warmup = mock.Mock(side_effect=error if is_fatal else [error, None]) + + with _no_cuda_side_effects() as (empty_cache, _synchronize): + if is_fatal: + with pytest.raises(torch.OutOfMemoryError) as excinfo: + engine._general_warmup_impl(object(), [(128, 0), (64, 0)]) + assert excinfo.value is error + else: + engine._general_warmup_impl(object(), [(128, 0), (64, 0)]) - indexer._project_and_quantize_q.assert_not_called() - synchronize.assert_not_called() - info.assert_not_called() - owner.dist.tp_allgather.assert_called_once_with(1) + if is_fatal: + # The remaining shape is never attempted: this rank's peers are + # already stuck in the failed forward's collectives. + engine._forward_warmup.assert_called_once() + else: + assert engine._forward_warmup.call_count == 2 + # A retry after an OOM between dispatch() and combine() has to start + # from a clean MoE all-to-all state. + engine._reset_moe_alltoall_state.assert_called_once_with() + empty_cache.assert_called_once_with() @pytest.mark.cpu_only -def test_indexer_q_prewarm_compiles_first_fused_indexer() -> None: - unfused_indexer = _deepseek_v4_indexer(fused_enabled=False) - fused_indexer = _deepseek_v4_indexer() - later_fused_indexer = _deepseek_v4_indexer() - owner = _indexer_q_prewarm_owner( - [unfused_indexer, fused_indexer, later_fused_indexer], tp_size=2 - ) - - with patch.object(torch.cuda, "synchronize"): - DecoderWarmup._prewarm_cute_dsl_indexer_q(owner) +@pytest.mark.parametrize( + "error,recoverable", + [ + (torch.OutOfMemoryError("OOM"), True), + (RuntimeError(_KV_ALLOC_ERROR), True), + (RuntimeError("unexpected"), False), + ], +) +def test_mamba_preforward_error_policy_when_alone(error: Exception, recoverable: bool) -> None: + engine, resource_manager = _mamba_engine() + engine._dummy_requests.create_warmup_request = mock.Mock(side_effect=error) + engine._forward_warmup = mock.Mock() + + with _no_cuda_side_effects(): + if recoverable: + _run_mamba_warmup(engine, resource_manager) + else: + with pytest.raises(RuntimeError, match="unexpected"): + _run_mamba_warmup(engine, resource_manager) - unfused_indexer._project_and_quantize_q.assert_not_called() - assert fused_indexer._project_and_quantize_q.call_count == 3 - later_fused_indexer._project_and_quantize_q.assert_not_called() - owner.dist.tp_allgather.assert_called_once_with(1) + engine._forward_warmup.assert_not_called() + # The failure predates dispatch(), so there is no half-finished MoE + # all-to-all exchange to unwind. + engine._reset_moe_alltoall_state.assert_not_called() @pytest.mark.cpu_only -def test_indexer_q_prewarm_failure_still_reaches_barrier() -> None: - """A per-shape failure must not skip the TP barrier the other ranks wait on.""" - indexer = _deepseek_v4_indexer() - indexer._project_and_quantize_q.side_effect = [RuntimeError("compile"), None, None] - owner = _indexer_q_prewarm_owner([indexer], tp_size=2) +def test_mamba_midforward_runtime_error_recovers_when_alone() -> None: + engine, resource_manager = _mamba_engine() + engine._dummy_requests.create_warmup_request = mock.Mock(return_value=object()) + engine._forward_warmup = mock.Mock(side_effect=RuntimeError("mid-forward failure")) - with ( - patch.object(torch.cuda, "synchronize"), - patch.object(decoder_warmup_module.logger, "warning") as warning, - ): - DecoderWarmup._prewarm_cute_dsl_indexer_q(owner) + with _no_cuda_side_effects(): + _run_mamba_warmup(engine, resource_manager) - assert indexer._project_and_quantize_q.call_count == 3 - warning.assert_called_once() - owner.dist.tp_allgather.assert_called_once_with(1) + # Every shape is attempted, and each failed forward has to leave the MoE + # all-to-all state clean for the shape that follows it. + attempted = engine._dummy_requests.create_warmup_request.call_count + assert attempted == 2 + assert engine._forward_warmup.call_count == attempted + assert engine._reset_moe_alltoall_state.call_count == attempted @pytest.mark.cpu_only -def test_indexer_q_prewarm_does_not_report_completion_when_every_compile_fails() -> None: - indexer = _deepseek_v4_indexer() - indexer._project_and_quantize_q.side_effect = RuntimeError("compile") - owner = _indexer_q_prewarm_owner([indexer], tp_size=2) +@pytest.mark.parametrize("phase", ["pre-forward", "mid-forward"]) +def test_mamba_error_is_fatal_when_distributed(phase: str) -> None: + engine, resource_manager = _mamba_engine(world_size=2) + error = RuntimeError(_KV_ALLOC_ERROR) + if phase == "pre-forward": + engine._dummy_requests.create_warmup_request = mock.Mock(side_effect=error) + engine._forward_warmup = mock.Mock() + else: + engine._dummy_requests.create_warmup_request = mock.Mock(return_value=object()) + engine._forward_warmup = mock.Mock(side_effect=error) - with ( - patch.object(torch.cuda, "synchronize"), - patch.object(decoder_warmup_module.logger, "warning") as warning, - patch.object(decoder_warmup_module.logger, "info") as info, - ): - DecoderWarmup._prewarm_cute_dsl_indexer_q(owner) + with _no_cuda_side_effects(): + with pytest.raises(RuntimeError) as excinfo: + _run_mamba_warmup(engine, resource_manager) - assert warning.call_count == 3 - info.assert_not_called() - owner.dist.tp_allgather.assert_called_once_with(1) + assert excinfo.value is error + # The remaining shape is never attempted: recovering locally would leave + # peers waiting in a forward this rank has abandoned. + engine._dummy_requests.create_warmup_request.assert_called_once() @pytest.mark.cpu_only -def test_indexer_q_prewarm_skips_indexers_of_other_pp_stages() -> None: - remote_indexer = _other_pp_stage_indexer() - local_indexer = _deepseek_v4_indexer() - owner = _indexer_q_prewarm_owner([remote_indexer, local_indexer], tp_size=2) - - with patch.object(torch.cuda, "synchronize"): - DecoderWarmup._prewarm_cute_dsl_indexer_q(owner) +class TestSamplingWarmup(unittest.TestCase): + def test_sampling_and_phase_metrics_without_kv_cache(self) -> None: + """Verify sampling warmup coverage and KV-cache cleanup timing within total warmup.""" + runner = object.__new__(DecoderRunner) + warmup = runner._warmup = decoder_warmup( + decoder_config(cuda_graph_batch_sizes=[1, 4]), + model=SimpleNamespace(config=SimpleNamespace(vocab_size=128)), + ) + resource_manager = Mock() + resource_manager.get_resource_manager.return_value = None + events = [] + phase = warmup._warmup_timer.phase - remote_indexer._project_and_quantize_q.assert_not_called() - assert local_indexer._project_and_quantize_q.call_count == 3 - owner.dist.tp_allgather.assert_called_once_with(1) + @contextlib.contextmanager + def record_metric_scope( + name: str, *, metrics: dict[str, float], metric_name: str + ) -> Iterator[None]: + events.append(("enter", metric_name)) + with phase(name, metrics=metrics, metric_name=metric_name): + yield + events.append(("exit", metric_name)) + with ( + patch.object( + warmup._warmup_timer, + "phase", + side_effect=record_metric_scope, + ), + patch( + "tensorrt_llm._torch.pyexecutor.engine.runners.decoder.warmup.warmup_sampling_module", + side_effect=lambda: events.append("sampling"), + ) as warmup_sampling, + _capture_tllm_logs() as logs, + ): + for enabled in (False, True): + with self.subTest(enable_in_graph_sampling=enabled): + events.clear() + warmup_sampling.reset_mock() + warmup._config = replace(warmup._config, enable_in_graph_sampling=enabled) + runner._eager_workspace_reclaimer = object() + runner.warmup(resource_manager) + self.assertEqual( + events, + [ + ("enter", "sampling_warmup_seconds"), + "sampling", + ("exit", "sampling_warmup_seconds"), + ], + ) + self.assertIsNone(runner._eager_workspace_reclaimer) + warmup_sampling.assert_called_once_with() -if __name__ == "__main__": - unittest.main() + self.assertTrue( + any("Skipping warm up as no KV Cache manager allocated." in log for log in logs) + ) + self.assertGreaterEqual(warmup._metrics["sampling_warmup_seconds"], 0) diff --git a/tests/unittest/_torch/executor/engine/runners/decoder/test_warmup_agreement.py b/tests/unittest/_torch/executor/engine/runners/decoder/test_warmup_agreement.py new file mode 100644 index 000000000000..7a7c988ce590 --- /dev/null +++ b/tests/unittest/_torch/executor/engine/runners/decoder/test_warmup_agreement.py @@ -0,0 +1,169 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +from _torch.executor.engine_test_utils import _agreement + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "world_size,dwdp_size,has_batch,expected", + [ + (1, 0, False, False), + # Every rank in the forward group lacks a batch. The allgather proves + # nobody is left inside the collective, so skipping strands no one -- + # this is the shape a max_batch_size=1 config simply cannot build. + (2, 0, False, False), + # DWDP peers live outside `self.dist`, so unanimity cannot be observed + # and the shape stays fatal. + (1, 4, False, None), + (2, 0, True, True), + ], +) +def test_warmup_batch_policy( + world_size: int, dwdp_size: int, has_batch: bool, expected: bool | None +) -> None: + agreement = _agreement(world_size=world_size, dwdp_size=dwdp_size) + batch = object() if has_batch else None + if expected is None: + with pytest.raises(RuntimeError, match="cannot skip the shape"): + agreement.should_run_warmup_batch(batch, 128, "general") + else: + assert agreement.should_run_warmup_batch(batch, 128, "general") is expected + + +@pytest.mark.cpu_only +def test_symmetric_absence_skips_instead_of_failing_the_world() -> None: + """The regression this fix exists for. + + `max_batch_size=1` cannot build the mixed context+generation attention + shape, on any rank, at any KV cache size. Every rank reaches the same + verdict from configuration alone, so the group agrees and skips. Failing + here took down ten multi-GPU tests that had always skipped this shape. + """ + agreement = _agreement(world_size=2, tp_size=2, peer_values=(0,)) + + assert ( + agreement.should_run_warmup_batch(None, 2, "attention, num_tokens=2, num_gen_requests=1") + is False + ) + + +@pytest.mark.cpu_only +def test_group_disagreement_still_fails_with_the_ranks_named() -> None: + """Asymmetric capacity is the deadlock the agreement exists to catch.""" + agreement = _agreement(world_size=3, peer_values=(1, 0)) + + with pytest.raises(RuntimeError) as excinfo: + agreement.should_run_warmup_batch(None, 128, "general") + + message = str(excinfo.value) + assert "rank(s) [0, 2]" in message + assert "would deadlock" in message + + +@pytest.mark.cpu_only +def test_pipeline_only_world_agrees_without_a_tp_group() -> None: + """TP=1 with pipeline peers is where the TP-only check gathered nothing. + + Three of the ten failures were disaggregated context servers at + `tensor_parallel_size: 1, pipeline_parallel_size: 2`. + """ + agreement = _agreement(world_size=2, tp_size=1, peer_values=(0,)) + + assert agreement.should_run_warmup_batch(None, 128, "general") is False + + +@pytest.mark.cpu_only +def test_agreement_group_is_absent_for_dwdp() -> None: + agreement = _agreement(world_size=4, dwdp_size=4) + + assert agreement.allgather() is None + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "world_size,peer_values,flag,expected", + [ + (1, None, True, True), + (2, (1,), True, True), + # One pipeline stage owns the guided decoder and opts out; the group + # follows it rather than letting the others enter alone. + (2, (0,), True, False), + (2, (1,), False, False), + ], +) +def test_phase_entry_flag_is_agreed( + world_size: int, peer_values: tuple[int, ...] | None, flag: bool, expected: bool +) -> None: + agreement = _agreement(world_size=world_size, peer_values=peer_values) + + assert agreement.agree_flag(flag) is expected + + +@pytest.mark.cpu_only +def test_shape_list_narrows_to_what_every_rank_proposed() -> None: + """Attention-DP gives ranks different capacities and so different shapes.""" + agreement = _agreement(world_size=2) + agreement.dist.allgather = lambda value: [value, [[1, 0], [2, 0]]] + + assert agreement.agree_shapes([(1, 0), (2, 0), (4096, 0)]) == [(1, 0), (2, 0)] + + +@pytest.mark.cpu_only +def test_shape_list_is_untouched_without_a_group() -> None: + agreement = _agreement(world_size=1) + configs = [(1, 0), (128, 0)] + + assert agreement.agree_shapes(configs) == configs + + +@pytest.mark.cpu_only +def test_warmup_batch_policy_without_communicator() -> None: + """An engine built without a communicator cannot strand anyone: skip, don't raise.""" + agreement = _agreement(with_dist=False) + + assert agreement.should_run_warmup_batch(None, 128, "general") is False + + +@pytest.mark.cpu_only +def test_tp_agreement_needs_a_communicator_not_just_a_mapping() -> None: + """A mapping can claim TP peers that no communicator exists to reach. + + The agreement check gates on `tp_size`, so it has to answer for that + state rather than walking into a collective on `None`. + """ + agreement = _agreement(tp_size=2, with_dist=False) + + assert agreement.should_run_warmup_batch(object(), 128, "general") is True + assert agreement.should_run_warmup_batch(None, 128, "general") is False + + +@pytest.mark.cpu_only +def test_tp_disagreement_reports_the_ranks_that_lost_their_batch() -> None: + """Asymmetric attention-DP capacity is the deadlock this check exists for. + + The rank that still has a batch must not walk into a forward its peer + will never join, and the failure has to name the peer so the KV cache + fraction can be raised. + + DWDP is what still reaches the TP-only check: its peers are outside + ``self.dist``, so no group agreement is available and the narrower TP + check remains the only thing standing between an asymmetric world and a + deadlock. + """ + agreement = _agreement(world_size=2, dwdp_size=4, tp_size=2, tp_peer_values=(0,)) + + with pytest.raises(RuntimeError) as excinfo: + agreement.should_run_warmup_batch(object(), 128, "general") + + message = str(excinfo.value) + assert "TP rank(s) [1]" in message + assert "[128, 0]" in message + + +@pytest.mark.cpu_only +def test_tp_agreement_lets_a_symmetric_world_run() -> None: + agreement = _agreement(world_size=2, tp_size=2, tp_peer_values=(1,)) + + assert agreement.should_run_warmup_batch(object(), 128, "general") is True diff --git a/tests/unittest/_torch/executor/engine/runners/decoder/utils.py b/tests/unittest/_torch/executor/engine/runners/decoder/utils.py new file mode 100644 index 000000000000..466ed871a13f --- /dev/null +++ b/tests/unittest/_torch/executor/engine/runners/decoder/utils.py @@ -0,0 +1,218 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + + +from __future__ import annotations + +from contextlib import nullcontext +from dataclasses import replace +from types import SimpleNamespace +from typing import TYPE_CHECKING +from unittest.mock import Mock + +import torch + +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention +from tensorrt_llm._torch.pyexecutor.engine.lora import LoraParamBuilder +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.config import DecoderRunnerConfig +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.dummy_requests import DummyRequestFactory +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.graphs import DecoderCudaGraphs +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.inputs import ( + DecoderInputBuffers, + DecoderInputPreparer, +) +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.warmup import ( + DecoderWarmup, + WarmupAgreement, +) +from tensorrt_llm._torch.pyexecutor.warmup_timer import _WarmupTimer +from tensorrt_llm.llmapi.llm_args import CudaGraphConfig, PrefillCudaGraphBackend +from tensorrt_llm.mapping import Mapping + +if TYPE_CHECKING: + from typing import Any + + from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import CUDAGraphRunner + from tensorrt_llm._torch.pyexecutor.engine.multimodal import MultimodalItemScheduler + from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.graphs import ForwardWarmup + from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.metadata import DecoderMetadata + + +def decoder_config(**changes: Any) -> DecoderRunnerConfig: + config = DecoderRunnerConfig( + max_batch_size=13, + max_num_tokens=258, + max_seq_len=2048, + max_beam_width=1, + without_logits=False, + attention_backend=TrtllmAttention, + dtype=torch.float16, + enable_attention_dp=False, + disable_overlap_scheduler=False, + is_spec_decode=False, + spec_config=None, + max_draft_len=0, + max_total_draft_tokens=0, + max_draft_loop_tokens=0, + original_max_draft_len=0, + original_max_total_draft_tokens=0, + spec_dec_max_total_draft_tokens=0, + num_seq_slots=None, + cuda_graph_config=CudaGraphConfig( + enable_padding=True, batch_sizes=[1, 2, 4, 8, 16, 32, 64, 128] + ), + cuda_graph_batch_sizes=[1, 2, 4, 8, 16, 32, 64, 128], + cuda_graph_padding_enabled=True, + max_cuda_graph_batch_size=128, + prefill_cuda_graph_backend=PrefillCudaGraphBackend.DISABLED, + prefill_cuda_graph_num_tokens=[], + enable_in_graph_sampling=False, + torch_compile_enabled=False, + torch_compile_piecewise_cuda_graph=False, + torch_compile_prefill_only=False, + compile_only_piecewise_graphs=False, + use_mrope=False, + is_multimodal=False, + mm_encoder_cache_enabled=False, + enable_autotuner=True, + cuda_graph_specialize_lora=False, + ) + spec = changes.get("spec_config") + if spec is not None and changes.get("is_spec_decode", True): + draft_len = spec.max_draft_len + total_draft_tokens = spec.tokens_per_gen_step - 1 + changes = dict( + { + "is_spec_decode": True, + "max_draft_len": draft_len, + "max_total_draft_tokens": total_draft_tokens, + "max_draft_loop_tokens": total_draft_tokens, + "original_max_draft_len": draft_len, + "original_max_total_draft_tokens": total_draft_tokens, + "spec_dec_max_total_draft_tokens": total_draft_tokens, + **changes, + } + ) + return replace(config, **changes) + + +def input_preparer( + config: DecoderRunnerConfig, + *, + model: torch.nn.Module | SimpleNamespace | None, + mapping: Mapping | None = None, + buffers: DecoderInputBuffers | None = None, + mm_item_scheduler: MultimodalItemScheduler | None = None, +) -> DecoderInputPreparer: + mapping = mapping or Mapping() + return DecoderInputPreparer( + buffers=buffers, + config=config, + mapping=mapping, + dist=None, + model=model, + get_runtime_tokens_per_gen_step=( + config.spec_config.get_runtime_tokens_per_gen_step + if config.is_spec_decode + else lambda draft_len: draft_len + 1 + ), + iter_states={}, + lora=LoraParamBuilder( + attn_backend=config.attention_backend, + cuda_graph_manager=None, + spec_config=config.spec_config, + ), + use_beam_search=config.max_beam_width > 1, + steady_gen_cache_supported=True, + mm_item_scheduler=mm_item_scheduler, + ) + + +def decoder_graphs( + config: DecoderRunnerConfig, + *, + model: torch.nn.Module | SimpleNamespace | None = None, + mapping: Mapping | None = None, + cuda_graph_runner: CUDAGraphRunner | None = None, + dummy_requests: DummyRequestFactory | None = None, + metadata: DecoderMetadata | None = None, + forward_warmup: ForwardWarmup | None = None, + lora: LoraParamBuilder | SimpleNamespace | None = None, +) -> DecoderCudaGraphs: + mapping = mapping or Mapping() + return DecoderCudaGraphs( + cuda_graph_runner=( + cuda_graph_runner + if cuda_graph_runner is not None + else Mock(enabled=False, is_warmup_only=False) + ), + breakable_cuda_graph_runner=None, + config=config, + mapping=mapping, + model=model if model is not None else torch.nn.Identity(), + sparse_attention_config=None, + metrics={}, + dummy_requests=dummy_requests if dummy_requests is not None else Mock(), + agreement=WarmupAgreement(None, mapping), + metadata=metadata if metadata is not None else Mock(), + get_runtime_tokens_per_gen_step=lambda draft_len: draft_len + 1, + dynamic_draft_len_mapping=None, + lora=lora if lora is not None else SimpleNamespace(cuda_graph_manager=None), + forward_warmup=forward_warmup if forward_warmup is not None else Mock(), + torch_compile_bypass=nullcontext, + context_promotion_supported=True, + ) + + +def decoder_warmup( + config: DecoderRunnerConfig, + *, + model: torch.nn.Module | SimpleNamespace | None = None, + mapping: Mapping | None = None, + agreement: WarmupAgreement | None = None, + dummy_requests: DummyRequestFactory | None = None, + graphs: DecoderCudaGraphs | None = None, + forward_warmup: ForwardWarmup | None = None, +) -> DecoderWarmup: + mapping = mapping or Mapping() + model = model if model is not None else torch.nn.Identity() + forward_warmup = forward_warmup if forward_warmup is not None else Mock() + dummy_requests = ( + dummy_requests + if dummy_requests is not None + else DummyRequestFactory( + config=config, + model=model, + get_runtime_tokens_per_gen_step=lambda draft_len: draft_len + 1, + release_request_state=None, + ) + ) + graphs = ( + graphs + if graphs is not None + else decoder_graphs( + config, + model=model, + mapping=mapping, + dummy_requests=dummy_requests, + forward_warmup=forward_warmup, + ) + ) + return DecoderWarmup( + config=config, + mapping=mapping, + dist=agreement.dist if agreement is not None else None, + model=model, + sparse_attention_config=None, + warmup_timer=_WarmupTimer(rank=0), + metrics={}, + agreement=agreement if agreement is not None else WarmupAgreement(None, mapping), + metadata=Mock(), + dummy_requests=dummy_requests, + graphs=graphs, + inputs=SimpleNamespace(guided_decoder=None), + forward_warmup=forward_warmup, + torch_compile_bypass=nullcontext, + context_warmups_supported=True, + eager_workspace_reclaimable=False, + ) diff --git a/tests/unittest/_torch/executor/engine/test_common.py b/tests/unittest/_torch/executor/engine/runners/test_common.py similarity index 79% rename from tests/unittest/_torch/executor/engine/test_common.py rename to tests/unittest/_torch/executor/engine/runners/test_common.py index 0134eaf6bfde..919783b06a06 100644 --- a/tests/unittest/_torch/executor/engine/test_common.py +++ b/tests/unittest/_torch/executor/engine/runners/test_common.py @@ -8,7 +8,13 @@ import torch from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata -from tensorrt_llm._torch.pyexecutor.engine.runners import common +from tensorrt_llm._torch.pyexecutor.engine.runners import ( + DecoderRunner, + EncoderDecoderRunner, + EncoderRunner, + common, + resolve_runner_type, +) from tensorrt_llm._torch.pyexecutor.engine.runners.common import ( apply_position_id_offset, get_all_rank_num_tokens, @@ -20,12 +26,73 @@ ship_multimodal_indices, token_model_inputs, ) +from tensorrt_llm._torch.pyexecutor.engine.runners.interface import ModelRunner +from tensorrt_llm._torch.pyexecutor.engine.runners.mm_encoder import MultimodalEncoderRunner +from tensorrt_llm._torch.pyexecutor.engine.runners.no_kv_cache import NoKVCacheRunner +from tensorrt_llm._torch.pyexecutor.engine.runners.pooling import PoolingRunner +from tensorrt_llm._torch.pyexecutor.model_engine import ModelEngine from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm.llmapi.llm_args import PrefillCudaGraphBackend pytestmark = pytest.mark.cpu_only +@pytest.mark.parametrize( + ("encode_only", "mm_encoder_only", "is_generation", "is_encoder_decoder", "runner_type"), + [ + (True, False, False, False, EncoderRunner), + (True, True, False, False, MultimodalEncoderRunner), + (False, True, False, False, MultimodalEncoderRunner), + (False, False, False, True, PoolingRunner), + (False, False, True, True, EncoderDecoderRunner), + (False, False, True, False, DecoderRunner), + ], +) +def test_resolve_runner_dispatches_startup_families( + encode_only: bool, + mm_encoder_only: bool, + is_generation: bool, + is_encoder_decoder: bool, + runner_type: type[ModelRunner], +) -> None: + args = SimpleNamespace(encode_only=encode_only, mm_encoder_only=mm_encoder_only) + model = SimpleNamespace( + model_config=SimpleNamespace( + is_generation=is_generation, is_encoder_decoder=is_encoder_decoder + ) + ) + + assert resolve_runner_type(model, args) is runner_type + + +@pytest.mark.parametrize( + ("runner_type", "initializer_name"), + [ + (EncoderRunner, "_initialize_encoder_runner"), + (EncoderDecoderRunner, "_initialize_encoder_decoder_runner"), + (NoKVCacheRunner, "_initialize_no_kv_cache_runner"), + (DecoderRunner, "_initialize_decoder_runner"), + (ModelRunner, None), + ], +) +def test_model_engine_initializes_runner_by_family( + runner_type: type[ModelRunner], initializer_name: str | None +) -> None: + engine = object.__new__(ModelEngine) + engine._model_caller = Mock() + if initializer_name is None: + with pytest.raises(TypeError, match="No runner initializer registered"): + engine._initialize_runner(runner_type) + return + + expected = Mock(spec=runner_type) + initializer = Mock(return_value=expected) + setattr(engine, initializer_name, initializer) + + assert engine._initialize_runner(runner_type) is expected + initializer.assert_called_once_with(runner_type) + + @pytest.mark.parametrize("token_id_attribute", ["multimodal_token_ids", "mm_token_ids"]) def test_prepare_multimodal_indices_uses_model_token_ids(token_id_attribute: str) -> None: model = SimpleNamespace(config=SimpleNamespace(vocab_size=100)) diff --git a/tests/unittest/_torch/executor/engine/test_encoder.py b/tests/unittest/_torch/executor/engine/runners/test_encoder.py similarity index 95% rename from tests/unittest/_torch/executor/engine/test_encoder.py rename to tests/unittest/_torch/executor/engine/runners/test_encoder.py index 1c90211e2746..a3dcafe2a117 100644 --- a/tests/unittest/_torch/executor/engine/test_encoder.py +++ b/tests/unittest/_torch/executor/engine/runners/test_encoder.py @@ -1,8 +1,12 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + from contextlib import contextmanager from types import SimpleNamespace +from typing import TYPE_CHECKING from unittest.mock import Mock, patch import pytest @@ -28,6 +32,11 @@ from tensorrt_llm._torch.pyexecutor.engine.runners.interface import PackedInputs from tensorrt_llm.llmapi.llm_args import EncodeCudaGraphConfig, EncodeExtraInputSpec +if TYPE_CHECKING: + from collections.abc import Callable, Iterator + + from tensorrt_llm._torch.attention.backends.interface import AttentionBackend + pytestmark = pytest.mark.cpu_only @@ -203,7 +212,7 @@ def test_multi_item_batch_stays_eager_even_when_a_graph_is_available() -> None: runner = object.__new__(EncoderRunner) @contextmanager - def pad_batch(inputs, batch_size): + def pad_batch(inputs: dict[str, object], batch_size: int) -> Iterator[dict[str, object]]: yield dict(inputs) graph_runner = SimpleNamespace( @@ -286,7 +295,9 @@ def test_encoder_runner_forwards_model_inputs_to_eager_preparation() -> None: ) -def _graph_enabled_runner(*, attention_backend=TrtllmAttention) -> EncoderRunner: +def _graph_enabled_runner( + *, attention_backend: type[AttentionBackend] = TrtllmAttention +) -> EncoderRunner: """An EncoderRunner whose graph runner is enabled but has no declared extra inputs.""" runner = object.__new__(EncoderRunner) runner._encoder_cuda_graph_runner = SimpleNamespace( @@ -369,7 +380,7 @@ def test_encoder_capture_runs_warmup_then_capture_and_restores_phase() -> None: events: list[bool] = [] @contextmanager - def allow_capture(): + def allow_capture() -> Iterator[None]: yield runner = object.__new__(EncoderRunner) @@ -394,7 +405,7 @@ def test_token_graph_capture_and_replay_enter_moe_iteration_context() -> None: capture_forward_calls: list[dict[str, object]] = [] @contextmanager - def moe_context(balance: object | None): + def moe_context(balance: object | None) -> Iterator[None]: context_balancers.append(balance) yield @@ -407,12 +418,16 @@ def needs_capture(key: tuple[int, int, int]) -> bool: return True @staticmethod - def capture(key, forward, inputs): + def capture( + key: tuple[int, int, int], + forward: Callable[[dict[str, object]], dict[str, torch.Tensor]], + inputs: dict[str, object], + ) -> dict[str, torch.Tensor]: capture_forward_calls.append(inputs) return forward(inputs) @staticmethod - def replay(key, inputs): + def replay(key: tuple[int, int, int], inputs: dict[str, object]) -> dict[str, torch.Tensor]: return {"logits": torch.tensor([[2.0]])} runner = object.__new__(EncoderRunner) diff --git a/tests/unittest/_torch/executor/engine/test_encoder_decoder.py b/tests/unittest/_torch/executor/engine/runners/test_encoder_decoder.py similarity index 99% rename from tests/unittest/_torch/executor/engine/test_encoder_decoder.py rename to tests/unittest/_torch/executor/engine/runners/test_encoder_decoder.py index ed949bbede69..a78762628b54 100644 --- a/tests/unittest/_torch/executor/engine/test_encoder_decoder.py +++ b/tests/unittest/_torch/executor/engine/runners/test_encoder_decoder.py @@ -1,6 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 + from contextlib import nullcontext from dataclasses import dataclass from types import SimpleNamespace @@ -201,7 +202,7 @@ def test_encoder_decoder_runner_builds_stage_before_decoder() -> None: order = Mock() - def decoder_init(self, *args, **kwargs) -> None: + def decoder_init(self: EncoderDecoderRunner, *args: object, **kwargs: object) -> None: order.decoder(self._encoder_graph_shapes, *args, **kwargs) with ( diff --git a/tests/unittest/_torch/executor/engine/runners/test_mm_encoder.py b/tests/unittest/_torch/executor/engine/runners/test_mm_encoder.py new file mode 100644 index 000000000000..f57e7776317a --- /dev/null +++ b/tests/unittest/_torch/executor/engine/runners/test_mm_encoder.py @@ -0,0 +1,206 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import json +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch +from _torch.executor.engine.runners.utils import _make_runner + +from tensorrt_llm._torch.pyexecutor.engine.runners.mm_encoder import MultimodalEncoderRunner +from tensorrt_llm._torch.pyexecutor.llm_request import PyResult +from tensorrt_llm._torch.pyexecutor.sampler import EarlyStopWithMMResult +from tensorrt_llm._torch.shared_tensor import SharedTensorContainer +from tensorrt_llm.inputs.multimodal import ( + MULTIMODAL_ENCODER_ITEM_METADATA_KEY, + MultimodalItemMetadata, +) + +pytestmark = pytest.mark.cpu_only + + +def test_mm_encoder_runner_forward_step_returns_empty_result_without_multimodal_params() -> None: + runner = _make_runner(MultimodalEncoderRunner, SimpleNamespace()) + + result = runner._forward_step({}, SimpleNamespace(context_requests=[])) + + assert result == { + "mm_embeddings": [], + "mm_embedding_request_indices": [], + "mm_embedding_lengths": [], + } + + +def test_mm_encoder_runner_forward_step_rejects_payload_request_count_mismatch() -> None: + request = SimpleNamespace(py_multimodal_data={"image": object()}) + runner = _make_runner(MultimodalEncoderRunner, SimpleNamespace()) + + with pytest.raises(ValueError, match="one multimodal payload per context"): + runner._forward_step( + {"multimodal_params": [SimpleNamespace(), SimpleNamespace()]}, + SimpleNamespace(context_requests=[request]), + ) + + +def test_mm_encoder_runner_forward_step_skips_metadata_only_and_missing_length_requests() -> None: + requests = [ + SimpleNamespace(py_multimodal_data={"mrope_config": {}}), + SimpleNamespace(py_multimodal_data={"image": object()}), + SimpleNamespace( + py_multimodal_data={ + "image": object(), + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( + [("image", 0)], output_embedding_lengths=[2] + ), + }, + multimodal_lengths=[2], + ), + ] + params = [ + SimpleNamespace(multimodal_data={}), + SimpleNamespace(multimodal_data={}), + SimpleNamespace(multimodal_data={}), + ] + model = SimpleNamespace(forward=Mock(return_value=[torch.arange(4).reshape(2, 2)])) + runner = _make_runner(MultimodalEncoderRunner, model) + + result = runner._forward_step( + {"multimodal_params": params}, + SimpleNamespace(context_requests=requests), + ) + + model.forward.assert_called_once_with([params[2]]) + assert result["logits"] is None + assert result["mm_embedding_request_indices"] == [2] + assert result["mm_embedding_lengths"] == [[2]] + torch.testing.assert_close(result["mm_embeddings"][0], torch.arange(4).reshape(2, 2)) + + +def test_mm_encoder_runner_splits_embeddings_and_returns_mrope_metadata() -> None: + first_param = SimpleNamespace( + multimodal_data={ + "image": object(), + "mrope_config": { + "mrope_position_ids": "first-ids", + "mrope_position_deltas": "first-deltas", + }, + } + ) + second_param = SimpleNamespace( + multimodal_data={ + "image": object(), + "mrope_config": { + "mrope_position_ids": "second-ids", + "mrope_position_deltas": "second-deltas", + }, + } + ) + requests = [ + SimpleNamespace( + py_multimodal_data={ + "image": object(), + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( + [("image", 0)], output_embedding_lengths=[1] + ), + }, + multimodal_lengths=[1], + ), + SimpleNamespace( + py_multimodal_data={ + "image": object(), + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( + [("image", 0)], output_embedding_lengths=[2] + ), + }, + multimodal_lengths=[2], + ), + ] + embeddings = torch.arange(6).reshape(3, 2) + model = SimpleNamespace(forward=Mock(return_value=[embeddings])) + runner = _make_runner(MultimodalEncoderRunner, model) + + result = runner._forward_step( + {"multimodal_params": [first_param, second_param]}, + SimpleNamespace(context_requests=requests), + ) + + model.forward.assert_called_once_with([first_param, second_param]) + assert result["logits"] is None + assert result["mm_embedding_request_indices"] == [0, 1] + assert result["mm_embedding_lengths"] == [[1], [2]] + torch.testing.assert_close(result["mm_embeddings"][0], embeddings[:1]) + torch.testing.assert_close(result["mm_embeddings"][1], embeddings[1:]) + assert result["mrope_position_ids"] == ["first-ids", "second-ids"] + assert result["mrope_position_deltas"] == ["first-deltas", "second-deltas"] + + +def test_mm_encoder_handoff_preserves_per_item_metadata_in_sparse_batch() -> None: + requests = [ + SimpleNamespace( + py_multimodal_data={ + "image": object(), + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( + [("image", 0)], output_embedding_lengths=[1] + ), + }, + multimodal_lengths=[1], + ), + SimpleNamespace(py_multimodal_data=None), + SimpleNamespace( + py_multimodal_data={ + "video": object(), + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( + [("video", 0), ("video", 1)], output_embedding_lengths=[3, 2] + ), + }, + multimodal_lengths=[7, 6], + ), + ] + for request in requests: + request.py_result = PyResult(prompt_len=1, max_new_tokens=1) + request.set_finished_reason = Mock() + params = [ + SimpleNamespace(multimodal_data=request.py_multimodal_data) + for request in requests + if request.py_multimodal_data is not None + ] + scheduled_requests = SimpleNamespace( + context_requests=requests, generation_requests=[], num_context_requests=len(requests) + ) + embeddings = torch.arange(12, dtype=torch.float32).reshape(6, 2) + + def encode(batch_params: list[SimpleNamespace]) -> list[torch.Tensor]: + batch_params[1].multimodal_data["multimodal_embedding_metadata"] = [ + {"retained_token_counts": [2, 1]}, + {"retained_token_counts": [2, 0]}, + ] + return [embeddings] + + runner = _make_runner(MultimodalEncoderRunner, SimpleNamespace(forward=encode)) + result = runner._forward_step( + {"multimodal_params": params}, + scheduled_requests, + ) + sampler = EarlyStopWithMMResult() + state = sampler.sample_async(scheduled_requests, result, []) + sampler.update_requests(state) + + assert requests[1].py_result.mm_embedding_handles is None + image_handles = json.loads(json.dumps(requests[0].py_result.mm_embedding_handles)) + assert len(image_handles) == 1 + assert "metadata" not in image_handles[0] + torch.testing.assert_close( + SharedTensorContainer.from_dict(image_handles[0]).get_local_view(), embeddings[:1] + ) + video_handles = json.loads(json.dumps(requests[2].py_result.mm_embedding_handles)) + assert [handle["metadata"]["retained_token_counts"] for handle in video_handles] == [ + [2, 1], + [2, 0], + ] + restored = [ + SharedTensorContainer.from_dict(handle).get_local_view() for handle in video_handles + ] + assert [len(embedding) for embedding in restored] == [3, 2] + torch.testing.assert_close(torch.cat(restored), embeddings[1:]) diff --git a/tests/unittest/_torch/executor/engine/test_no_kv_cache.py b/tests/unittest/_torch/executor/engine/runners/test_no_kv_cache.py similarity index 65% rename from tests/unittest/_torch/executor/engine/test_no_kv_cache.py rename to tests/unittest/_torch/executor/engine/runners/test_no_kv_cache.py index f76f0a97f577..78bbf26cb338 100644 --- a/tests/unittest/_torch/executor/engine/test_no_kv_cache.py +++ b/tests/unittest/_torch/executor/engine/runners/test_no_kv_cache.py @@ -1,18 +1,31 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 + from types import SimpleNamespace from unittest.mock import Mock, patch import pytest import torch +from _torch.executor.engine.runners.utils import _config, _make_runner, _model from tensorrt_llm._torch.attention.backends.interface import AttentionRuntimeFeatures from tensorrt_llm._torch.pyexecutor.engine.input_buffers import InputBuffers from tensorrt_llm._torch.pyexecutor.engine.runners import common as common_module from tensorrt_llm._torch.pyexecutor.engine.runners import no_kv_cache as no_kv_cache_module -from tensorrt_llm._torch.pyexecutor.engine.runners.no_kv_cache import NoKVCacheRunnerConfig +from tensorrt_llm._torch.pyexecutor.engine.runners.interface import ( + PreparedInputs, + RunnerConfig, + ScheduledInputs, +) +from tensorrt_llm._torch.pyexecutor.engine.runners.mm_encoder import MultimodalEncoderRunner +from tensorrt_llm._torch.pyexecutor.engine.runners.no_kv_cache import ( + NoKVCacheRunner, + NoKVCacheRunnerConfig, +) from tensorrt_llm._torch.pyexecutor.engine.runners.pooling import PoolingRunner +from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests +from tensorrt_llm._torch.utils import get_model_extra_attrs from tensorrt_llm.inputs.multimodal import MultimodalParams from tensorrt_llm.llmapi.llm_args import PrefillCudaGraphBackend @@ -226,3 +239,97 @@ def test_no_kv_cache_runner_prepare_inputs_gathers_attention_dp_token_counts( assert attn_metadata.all_rank_num_tokens == [2, 5] dist.tp_allgather_int64.assert_called_once_with([2]) dist.tp_cp_allgather_int64.assert_called_once_with([2]) + + +def test_no_kv_cache_runner_rejects_kv_manager_before_preparation() -> None: + runner = _make_runner(PoolingRunner, _model(is_generation=False)) + runner._prepare_inputs = Mock() + resource_manager = Mock() + resource_manager.get_resource_manager.return_value = object() + + with pytest.raises( + AssertionError, + match="no-KV-cache runner was initialized, but a KV cache manager was allocated", + ): + runner.forward( + ScheduledInputs(batch=ScheduledRequests()), resource_manager=resource_manager + ) + + runner._prepare_inputs.assert_not_called() + + +def test_no_kv_cache_config_extends_common_runner_config() -> None: + assert isinstance(_config(), RunnerConfig) + + +def test_no_kv_cache_runner_exposes_only_forward_step_as_abstract() -> None: + assert NoKVCacheRunner.__abstractmethods__ == frozenset({"_forward_step"}) + + +@pytest.mark.parametrize("runner_type", [PoolingRunner, MultimodalEncoderRunner]) +def test_concrete_runners_implement_no_kv_cache_forward_step( + runner_type: type[NoKVCacheRunner], +) -> None: + assert issubclass(runner_type, NoKVCacheRunner) + assert runner_type._forward_step is not NoKVCacheRunner._forward_step + + +def test_no_kv_cache_runner_owns_and_reuses_attention_metadata() -> None: + runner = _make_runner(PoolingRunner, _model(is_generation=False)) + + first = runner._setup_attn_metadata() + second = runner._setup_attn_metadata() + + assert first is second + assert first.kv_cache_manager is None + assert first.max_num_requests == 4 + assert first.max_num_sequences == 4 + assert first.max_num_heads_per_kv == 4 + assert first.block_ids_per_seq is None + assert first.kv_block_ids_per_seq is None + + +@pytest.mark.parametrize("runner_type", [PoolingRunner, MultimodalEncoderRunner]) +def test_no_kv_cache_runner_default_lifecycle_with_no_kv_resources( + runner_type: type[PoolingRunner] | type[MultimodalEncoderRunner], +) -> None: + runner = _make_runner(runner_type, _model(is_generation=False)) + resource_manager = SimpleNamespace(get_resource_manager=Mock(return_value=None)) + + assert runner.warmup(resource_manager) is None + assert runner.release_graphs() is None + assert runner.wait_for_input_copy() is None + + +def test_no_kv_forward_enters_model_context() -> None: + model = _model(is_generation=False) + + def model_forward(**kwargs: object) -> dict[str, object]: + assert torch.is_inference_mode_enabled() + assert get_model_extra_attrs() is model.extra_attrs + return {"hidden_states": object()} + + caller = Mock(side_effect=model_forward) + runner = _make_runner(PoolingRunner, model, model_caller=caller) + runner._prepare_inputs = Mock(return_value=PreparedInputs({})) + resources = SimpleNamespace(get_resource_manager=Mock(return_value=None)) + + runner.forward(ScheduledInputs(batch=ScheduledRequests()), resource_manager=resources) + + caller.assert_called_once() + assert get_model_extra_attrs() is None + + +def test_no_kv_forward_preserves_model_output_dictionary_without_length_update() -> None: + outputs = {"hidden_states": object()} + caller = Mock(return_value=outputs) + runner = _make_runner(PoolingRunner, _model(is_generation=False), model_caller=caller) + runner._prepare_inputs = Mock(return_value=PreparedInputs({})) + inputs = ScheduledInputs(batch=ScheduledRequests(), runtime_draft_len=3) + resources = SimpleNamespace(get_resource_manager=Mock(return_value=None)) + + actual = runner.forward(inputs, resource_manager=resources) + + assert actual is outputs + assert "runtime_draft_len" not in outputs + assert inputs.runtime_draft_len == 3 diff --git a/tests/unittest/_torch/executor/engine/runners/test_pooling.py b/tests/unittest/_torch/executor/engine/runners/test_pooling.py new file mode 100644 index 000000000000..56568c897d9b --- /dev/null +++ b/tests/unittest/_torch/executor/engine/runners/test_pooling.py @@ -0,0 +1,61 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from dataclasses import replace +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch +from _torch.executor.engine.runners.utils import _config, _make_runner, _model + +from tensorrt_llm._torch.pyexecutor.engine.runners.interface import PreparedInputs, ScheduledInputs +from tensorrt_llm._torch.pyexecutor.engine.runners.pooling import PoolingRunner +from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests + +pytestmark = pytest.mark.cpu_only + + +def test_pooling_runner_owns_forward_output_processing( + monkeypatch: pytest.MonkeyPatch, +) -> None: + logits = torch.arange(12).reshape(3, 4) + model_forward = Mock(return_value=logits) + runner = _make_runner( + PoolingRunner, + _model(is_generation=False), + model_caller=model_forward, + ) + attn_metadata = SimpleNamespace(on_update_kv_lens=Mock()) + prepared = PreparedInputs( + {"attn_metadata": attn_metadata, "input_ids": torch.tensor([1])}, + ) + monkeypatch.setattr(runner, "_prepare_inputs", Mock(return_value=prepared)) + + outputs = runner.forward( + ScheduledInputs(batch=ScheduledRequests()), + resource_manager=SimpleNamespace(get_resource_manager=Mock(return_value=None)), + ) + + torch.testing.assert_close(outputs["logits"], logits) + attn_metadata.on_update_kv_lens.assert_called_once_with() + model_forward.assert_called_once_with( + attn_metadata=attn_metadata, + input_ids=prepared.kwargs["input_ids"], + return_context_logits=False, + ) + + +def test_pooling_runner_returns_raw_model_outputs_when_logits_are_disabled() -> None: + model_outputs = {"hidden_states": object()} + model_forward = Mock(return_value=model_outputs) + runner = _make_runner( + PoolingRunner, + _model(is_generation=False), + replace(_config(), without_logits=True), + model_caller=model_forward, + ) + + outputs = runner._forward_step({}, SimpleNamespace()) + + assert outputs is model_outputs diff --git a/tests/unittest/_torch/executor/engine/runners/utils.py b/tests/unittest/_torch/executor/engine/runners/utils.py new file mode 100644 index 000000000000..d06539a5a8ab --- /dev/null +++ b/tests/unittest/_torch/executor/engine/runners/utils.py @@ -0,0 +1,73 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + + +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import torch +from _torch.executor.engine_test_utils import _AttentionBackend + +from tensorrt_llm._torch.attention.backends.interface import AttentionRuntimeFeatures +from tensorrt_llm._torch.pyexecutor.engine.input_buffers import InputBuffers +from tensorrt_llm._torch.pyexecutor.engine.runners.mm_encoder import MultimodalEncoderRunner +from tensorrt_llm._torch.pyexecutor.engine.runners.no_kv_cache import NoKVCacheRunnerConfig +from tensorrt_llm._torch.pyexecutor.engine.runners.pooling import PoolingRunner +from tensorrt_llm.llmapi.llm_args import PrefillCudaGraphBackend + + +def _make_runner( + runner_type: type[PoolingRunner] | type[MultimodalEncoderRunner], + model: torch.nn.Module | SimpleNamespace, + config: NoKVCacheRunnerConfig | None = None, + *, + model_caller: Mock | None = None, +) -> PoolingRunner | MultimodalEncoderRunner: + config = config or _config() + buffers = InputBuffers( + input_ids_cuda=torch.empty(config.max_num_tokens, dtype=torch.int), + position_ids_cuda=torch.empty(config.max_num_tokens, dtype=torch.int), + ) + kwargs = dict( + mapping=SimpleNamespace(), + dist=SimpleNamespace(), + moe_load_balancer=None, + attention_runtime_features=AttentionRuntimeFeatures(), + ) + if issubclass(runner_type, PoolingRunner): + kwargs["model_caller"] = model_caller or Mock(return_value={"logits": torch.empty(0)}) + with patch.object(InputBuffers, "allocate", return_value=buffers): + return runner_type(model, config, **kwargs) + + +def _config() -> NoKVCacheRunnerConfig: + return NoKVCacheRunnerConfig( + max_batch_size=4, + max_num_tokens=16, + max_seq_len=16, + max_beam_width=1, + without_logits=False, + attention_backend=_AttentionBackend, + enable_attention_dp=False, + prefill_cuda_graph_backend=PrefillCudaGraphBackend.DISABLED, + prefill_cuda_graph_num_tokens=[], + mm_encoder_cache_enabled=False, + ) + + +def _model(*, is_generation: bool, is_encoder_decoder: bool = False) -> SimpleNamespace: + pretrained_config = SimpleNamespace( + num_attention_heads=8, + num_key_value_heads=2, + ) + return SimpleNamespace( + extra_attrs={}, + model_config=SimpleNamespace( + is_generation=is_generation, + is_encoder_decoder=is_encoder_decoder, + pretrained_config=pretrained_config, + sparse_attention_config=None, + enable_flash_mla=False, + ), + config=pretrained_config, + ) diff --git a/tests/unittest/_torch/executor/engine/test_input_buffers.py b/tests/unittest/_torch/executor/engine/test_input_buffers.py new file mode 100644 index 000000000000..fa1c3199f77c --- /dev/null +++ b/tests/unittest/_torch/executor/engine/test_input_buffers.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + + +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor.engine.input_buffers import InputBuffers + + +@pytest.mark.parametrize("draft_capacity", [None, 0, 8]) +def test_runner_input_buffers_support_repeated_cuda_graph_replay( + draft_capacity: int | None, +) -> None: + buffers = InputBuffers.allocate( + max_num_tokens=16, + max_batch_size=2, + max_beam_width=2, + max_seq_len=8, + use_cache_indirection=True, + max_num_draft_tokens=draft_capacity, + ) + assert buffers.cache_indirection.shape == (2, 2, 8) + if draft_capacity is None: + assert buffers.gather_ids_cuda is None + assert buffers.draft_tokens_cuda is None + else: + assert buffers.gather_ids_cuda.numel() == 16 + assert buffers.draft_tokens_cuda.numel() == draft_capacity + input_ptr = buffers.input_ids_cuda.data_ptr() + buffers.input_ids_cuda.fill_(1) + buffers.position_ids_cuda.fill_(2) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + output = buffers.input_ids_cuda + buffers.position_ids_cuda + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = buffers.input_ids_cuda + buffers.position_ids_cuda + output_ptr = output.data_ptr() + for token in (3, 7): + buffers.input_ids_cuda.fill_(token) + graph.replay() + torch.testing.assert_close(output, torch.full_like(output, token + 2)) + assert output.data_ptr() == output_ptr + assert buffers.input_ids_cuda.data_ptr() == input_ptr + graph.reset() diff --git a/tests/unittest/_torch/executor/engine/test_model_call.py b/tests/unittest/_torch/executor/engine/test_model_call.py new file mode 100644 index 000000000000..715cd076072a --- /dev/null +++ b/tests/unittest/_torch/executor/engine/test_model_call.py @@ -0,0 +1,76 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + + +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import pytest +from _torch.executor.engine_test_utils import _AttentionMetadata + +from tensorrt_llm._torch.compilation.backend import Backend +from tensorrt_llm._torch.pyexecutor.engine.model_call import ModelCaller +from tensorrt_llm._torch.utils import get_model_extra_attrs, model_extra_attrs + + +@pytest.mark.cpu_only +def test_model_caller_uses_current_forward_and_restores_outer_attribute_context() -> None: + metadata = _AttentionMetadata() + spec_metadata = object() + attrs = {} + outer_attrs = {"outer": True} + model = SimpleNamespace(model_config=SimpleNamespace(extra_attrs={"model_flag": True})) + + def first_forward(**kwargs: object) -> str: + assert get_model_extra_attrs() is attrs + assert attrs["attention_metadata"]() is metadata + assert attrs["spec_metadata"] is spec_metadata + assert attrs["model_flag"] is True + return "first" + + model.forward = Mock(side_effect=first_forward) + caller = ModelCaller(model) + with model_extra_attrs(outer_attrs): + with model_extra_attrs(attrs): + assert caller(attn_metadata=metadata, spec_metadata=spec_metadata) == "first" + model.forward.assert_called_once_with( + attn_metadata=metadata, spec_metadata=spec_metadata + ) + model.forward = Mock(return_value="replacement") + assert caller(attn_metadata=metadata) == "replacement" + model.forward.assert_called_once_with(attn_metadata=metadata) + assert attrs["spec_metadata"] is None + model.forward.side_effect = RuntimeError("model failure") + with pytest.raises(RuntimeError, match="model failure"): + caller(attn_metadata=metadata) + assert get_model_extra_attrs() is outer_attrs + + +@pytest.mark.cpu_only +def test_model_caller_borrows_live_compile_streams_and_events() -> None: + streams = Backend.Streams() + backend = SimpleNamespace(events=Backend.Events()) + metadata = _AttentionMetadata() + current_stream = object() + attrs = {} + + def forward(**kwargs: object) -> tuple[int, int]: + assert get_model_extra_attrs() is attrs + assert attrs["aux_streams"]() is streams + assert attrs["events"]() is backend.events + assert attrs["global_stream"] is current_stream + return len(attrs["aux_streams"]()), len(attrs["events"]()) + + model = SimpleNamespace(model_config=SimpleNamespace(extra_attrs={}), forward=forward) + caller = ModelCaller(model, compile_backend=backend, aux_streams=streams) + with patch("torch.cuda.current_stream", return_value=current_stream), model_extra_attrs(attrs): + assert caller(attn_metadata=metadata) == (0, 0) + streams.append(object()) + backend.events = Backend.Events([object(), object()]) + assert caller(attn_metadata=metadata) == (1, 2) + + +@pytest.mark.cpu_only +def test_model_caller_requires_streams_with_compile_backend() -> None: + with pytest.raises(ValueError, match="requires its auxiliary stream container"): + ModelCaller(SimpleNamespace(), compile_backend=SimpleNamespace(events=Backend.Events())) diff --git a/tests/unittest/_torch/executor/engine/test_multimodal.py b/tests/unittest/_torch/executor/engine/test_multimodal.py index 7abf27d0e11d..ea35b7f0e6fc 100644 --- a/tests/unittest/_torch/executor/engine/test_multimodal.py +++ b/tests/unittest/_torch/executor/engine/test_multimodal.py @@ -1,7 +1,12 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import unittest from types import SimpleNamespace +from typing import TYPE_CHECKING from unittest.mock import Mock, patch import pytest @@ -12,10 +17,12 @@ make_mm_request, ) +import tensorrt_llm from tensorrt_llm._torch.models.modeling_multimodal_encoder import MultimodalEncoderMixin from tensorrt_llm._torch.models.modeling_multimodal_mixin import ( MultimodalEncoderContractError, MultimodalModelMixin, + _build_request_multimodal_input, ) from tensorrt_llm._torch.pyexecutor.engine.multimodal import ( MultimodalItemScheduler, @@ -44,12 +51,17 @@ MultimodalEncoderCudaGraphConfig, MultimodalEncoderSchedulingPolicy, ) +from tensorrt_llm.mapping import Mapping + +if TYPE_CHECKING: + from collections.abc import Sequence + from types import TracebackType # The item-scheduling surface is pure logic: no kernels, no device transfers. pytestmark = pytest.mark.cpu_only -def test_disabled_encoder_does_not_initialize_processor(): +def test_disabled_encoder_does_not_initialize_processor() -> None: processor = Mock(spec=BaseMultimodalDummyInputsBuilder) processor.get_mm_max_tokens_per_item.side_effect = AssertionError("encoder is disabled") model = torch.nn.Module() @@ -224,7 +236,6 @@ def test_engine_captures_encoder_graphs_with_original_token_budget( item_scheduling: bool, ) -> None: from tensorrt_llm._torch.pyexecutor import _util, model_engine - from tensorrt_llm.mapping import Mapping class Encoder(torch.nn.Module, MultimodalEncoderMixin): encoder_cuda_graph_modality = "vision" @@ -332,7 +343,9 @@ def _build_blocks_graph_runner(self, config: MultimodalEncoderCudaGraphConfig) - @pytest.mark.parametrize("enable_padding", [False, True]) @pytest.mark.parametrize("automatic", [False, True], ids=["explicit", "automatic"]) -def test_encoder_graphs_preserve_eager_metadata_capacity(enable_padding, automatic) -> None: +def test_encoder_graphs_preserve_eager_metadata_capacity( + enable_padding: bool, automatic: bool +) -> None: class Encoder(torch.nn.Module, MultimodalEncoderMixin): def __init__(self) -> None: super().__init__() @@ -368,7 +381,9 @@ def test_encoder_graph_modality_validation_allows_delegated_encoders() -> None: ) -def _bind_items(mm_item_scheduler: MultimodalItemScheduler, request, *, row_bytes: int = 8) -> None: +def _bind_items( + mm_item_scheduler: MultimodalItemScheduler, request: LlmRequest, *, row_bytes: int = 8 +) -> None: state = request.py_mm_encoder_state for item_idx, rows in enumerate( request.py_multimodal_params.item_metadata.output_embedding_lengths @@ -504,11 +519,15 @@ def test_item_encoder_classifies_request_state_contract_errors() -> None: def test_item_encoder_classifies_output_count_contract_error() -> None: class _Model(MultimodalModelMixin): - def prepare_multimodal_encoder_inputs(self, _): + def prepare_multimodal_encoder_inputs( + self, _: Sequence[tuple[MultimodalParams, int]] + ) -> list[tuple[SimpleNamespace, list[int], str]]: encoder_input = SimpleNamespace(to_device=lambda *_args, **_kwargs: None) return [(encoder_input, [1], "image")] - def forward_multimodal_encoder_items(self, _): + def forward_multimodal_encoder_items( + self, _: Sequence[tuple[MultimodalParams, Sequence[int], str]] + ) -> list[torch.Tensor]: return [] mm_item_scheduler = bare_mm_item_scheduler(_Model()) @@ -522,13 +541,17 @@ def forward_multimodal_encoder_items(self, _): @pytest.mark.parametrize("failure_stage", ["prepare", "forward"]) def test_item_encoder_translates_model_contract_errors(failure_stage: str) -> None: class _Model(MultimodalModelMixin): - def prepare_multimodal_encoder_inputs(self, _): + def prepare_multimodal_encoder_inputs( + self, _: Sequence[tuple[MultimodalParams, int]] + ) -> list[tuple[SimpleNamespace, list[int], str]]: if failure_stage == "prepare": raise MultimodalEncoderContractError("bad request metadata") encoder_input = SimpleNamespace(to_device=lambda *_args, **_kwargs: None) return [(encoder_input, [1], "image")] - def forward_multimodal_encoder_items(self, _): + def forward_multimodal_encoder_items( + self, _: Sequence[tuple[MultimodalParams, Sequence[int], str]] + ) -> list[torch.Tensor]: raise MultimodalEncoderContractError("bad encoder output rows") mm_item_scheduler = bare_mm_item_scheduler(_Model()) @@ -550,13 +573,17 @@ def forward_multimodal_encoder_items(self, _): @pytest.mark.parametrize("failure_stage", ["prepare", "forward"]) def test_item_encoder_does_not_translate_system_errors(failure_stage: str) -> None: class _Model(MultimodalModelMixin): - def prepare_multimodal_encoder_inputs(self, _): + def prepare_multimodal_encoder_inputs( + self, _: Sequence[tuple[MultimodalParams, int]] + ) -> list[tuple[SimpleNamespace, list[int], str]]: if failure_stage == "prepare": raise torch.cuda.OutOfMemoryError("encoder OOM") encoder_input = SimpleNamespace(to_device=lambda *_args, **_kwargs: None) return [(encoder_input, [1], "image")] - def forward_multimodal_encoder_items(self, _): + def forward_multimodal_encoder_items( + self, _: Sequence[tuple[MultimodalParams, Sequence[int], str]] + ) -> list[torch.Tensor]: raise torch.cuda.OutOfMemoryError("encoder OOM") mm_item_scheduler = bare_mm_item_scheduler(_Model()) @@ -579,11 +606,15 @@ def test_item_encoder_closes_window_scope_before_request_errors( forward_failure: str, expected_error: type, scope_failed: bool ) -> None: class _Model(MultimodalModelMixin): - def prepare_multimodal_encoder_inputs(self, _): + def prepare_multimodal_encoder_inputs( + self, _: Sequence[tuple[MultimodalParams, int]] + ) -> list[tuple[SimpleNamespace, list[int], str]]: encoder_input = SimpleNamespace(to_device=lambda *_args, **_kwargs: None) return [(encoder_input, [1], "image")] - def forward_multimodal_encoder_items(self, _): + def forward_multimodal_encoder_items( + self, _: Sequence[tuple[MultimodalParams, Sequence[int], str]] + ) -> list[torch.Tensor]: if forward_failure == "contract": raise MultimodalEncoderContractError("bad encoder output rows") if forward_failure == "system": @@ -593,13 +624,18 @@ def forward_multimodal_encoder_items(self, _): class _RecordingScope: exits = [] - def __init__(self, _inputs): + def __init__(self, _inputs: object) -> None: pass - def __enter__(self): + def __enter__(self) -> _RecordingScope: return self - def __exit__(self, exc_type, _exc, _tb): + def __exit__( + self, + exc_type: type[BaseException] | None, + _exc: BaseException | None, + _tb: TracebackType | None, + ) -> bool: self.exits.append(exc_type is not None) return False @@ -624,10 +660,14 @@ def test_item_encoder_validates_all_outputs_before_commit(invalid_output: str) - class _Model(MultimodalModelMixin): embedding_dtype = torch.bfloat16 - def prepare_multimodal_encoder_inputs(self, _): + def prepare_multimodal_encoder_inputs( + self, _: Sequence[tuple[MultimodalParams, int]] + ) -> list[tuple[SimpleNamespace, list[int], str]]: return [] - def forward_multimodal_encoder_items(self, _): + def forward_multimodal_encoder_items( + self, _: Sequence[tuple[MultimodalParams, Sequence[int], str]] + ) -> list[torch.Tensor | None]: return outputs outputs = [torch.ones(1, 4, dtype=torch.bfloat16) for _ in range(3)] @@ -686,7 +726,9 @@ def test_item_outputs_commit_to_prompt_ordered_cache_keys( class _Model(MultimodalModelMixin): embedding_dtype = torch.float32 - def forward_multimodal_encoder_items(self, encoder_inputs): + def forward_multimodal_encoder_items( + self, encoder_inputs: Sequence[tuple[MultimodalParams, Sequence[int], str]] + ) -> list[torch.Tensor]: return [ torch.full((embedding_length, 2), float(embedding_length)) for _, embedding_lengths, _ in encoder_inputs @@ -819,3 +861,59 @@ class _Model(MultimodalModelMixin): torch.testing.assert_close(multimodal_data["multimodal_embedding"], torch.ones(2, 2)) # Only the per-forward payload drops the raw inputs. assert "image" in request.py_multimodal_data + + +class MultimodalSetupTestCase(unittest.TestCase): + def test_build_request_multimodal_input_skips_when_cache_disabled(self) -> None: + request = LlmRequest( + request_id=1, + max_new_tokens=1, + input_tokens=[0, 1, 2], + sampling_config=tensorrt_llm.bindings.SamplingConfig(1), + is_streaming=False, + multimodal_hashes=[[1, 2, 3, 4, 5, 6, 7, 8]], + multimodal_positions=[1], + multimodal_lengths=[1], + multimodal_uuids=["image-0"], + ) + + # With the encoder cache disabled, nothing consumes `multimodal_input`, + # so it should not be built at all. + self.assertIsNone(_build_request_multimodal_input(request, cache_enabled=False)) + + def test_multimodal_encoder_max_seq_len(self) -> None: + class CapturingEncoder(torch.nn.Module, MultimodalEncoderMixin): + def __init__(self) -> None: + super().__init__() + self.setup_args = None + self.max_seq_len = None + + def setup_attn_metadata(self, max_num_tokens: int) -> None: + self.setup_args = max_num_tokens + + def set_attn_max_seq_len(self, max_seq_len: int) -> None: + self.max_seq_len = max_seq_len + + encoder_max_num_tokens = 16384 + processor_max_num_tokens = 65536 + cases = [ + ({"image": processor_max_num_tokens}, processor_max_num_tokens), + ({"image": 4096}, encoder_max_num_tokens), + ({}, encoder_max_num_tokens), + (None, encoder_max_num_tokens), + ] + for max_tokens_per_item, expected_max_seq_len in cases: + with self.subTest(max_tokens_per_item=max_tokens_per_item): + encoder = CapturingEncoder() + if max_tokens_per_item is None: + input_processor = Mock() + else: + input_processor = Mock(spec=BaseMultimodalDummyInputsBuilder) + input_processor.get_mm_max_tokens_per_item.return_value = max_tokens_per_item + + setup_mm_encoder_attn_metadata( + torch.nn.Sequential(encoder), input_processor, encoder_max_num_tokens, None + ) + + self.assertEqual(encoder.setup_args, encoder_max_num_tokens) + self.assertEqual(encoder.max_seq_len, expected_max_seq_len) diff --git a/tests/unittest/_torch/executor/engine/test_runners.py b/tests/unittest/_torch/executor/engine/test_runners.py deleted file mode 100644 index ef51370f0bf7..000000000000 --- a/tests/unittest/_torch/executor/engine/test_runners.py +++ /dev/null @@ -1,819 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -import json -from dataclasses import FrozenInstanceError, replace -from types import SimpleNamespace -from typing import Any -from unittest.mock import Mock, patch - -import pytest -import torch - -from tensorrt_llm._torch.attention.backends.interface import AttentionRuntimeFeatures -from tensorrt_llm._torch.compilation.backend import Backend -from tensorrt_llm._torch.pyexecutor.engine.input_buffers import InputBuffers -from tensorrt_llm._torch.pyexecutor.engine.model_call import ModelCaller -from tensorrt_llm._torch.pyexecutor.engine.runners import resolve_runner_type -from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import DecoderRunner -from tensorrt_llm._torch.pyexecutor.engine.runners.encoder import EncoderRunner -from tensorrt_llm._torch.pyexecutor.engine.runners.encoder_decoder import EncoderDecoderRunner -from tensorrt_llm._torch.pyexecutor.engine.runners.interface import ( - ModelRunner, - PreparedInputs, - RunnerConfig, - ScheduledInputs, -) -from tensorrt_llm._torch.pyexecutor.engine.runners.mm_encoder import MultimodalEncoderRunner -from tensorrt_llm._torch.pyexecutor.engine.runners.no_kv_cache import ( - NoKVCacheRunner, - NoKVCacheRunnerConfig, -) -from tensorrt_llm._torch.pyexecutor.engine.runners.pooling import PoolingRunner -from tensorrt_llm._torch.pyexecutor.llm_request import PyResult -from tensorrt_llm._torch.pyexecutor.model_engine import ModelEngine -from tensorrt_llm._torch.pyexecutor.sampler import EarlyStopWithMMResult -from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests -from tensorrt_llm._torch.shared_tensor import SharedTensorContainer -from tensorrt_llm._torch.utils import get_model_extra_attrs, model_extra_attrs -from tensorrt_llm.inputs.multimodal import ( - MULTIMODAL_ENCODER_ITEM_METADATA_KEY, - MultimodalItemMetadata, -) -from tensorrt_llm.llmapi.llm_args import ( - CudaGraphConfig, - EncodeCudaGraphConfig, - PrefillCudaGraphBackend, -) - -pytestmark = pytest.mark.cpu_only - - -class _AttentionMetadata: - def __init__(self, **kwargs: Any) -> None: - self.__dict__.update(kwargs) - self.block_ids_per_seq = object() - self.kv_block_ids_per_seq = object() - self.on_update_kv_lens = Mock() - - -class _AttentionBackend: - Metadata = _AttentionMetadata - - -def _make_runner( - runner_type: type[PoolingRunner] | type[MultimodalEncoderRunner], - model: Any, - config: NoKVCacheRunnerConfig | None = None, - *, - model_caller: Mock | None = None, -) -> PoolingRunner | MultimodalEncoderRunner: - config = config or _config() - buffers = InputBuffers( - input_ids_cuda=torch.empty(config.max_num_tokens, dtype=torch.int), - position_ids_cuda=torch.empty(config.max_num_tokens, dtype=torch.int), - ) - kwargs = dict( - mapping=SimpleNamespace(), - dist=SimpleNamespace(), - moe_load_balancer=None, - attention_runtime_features=AttentionRuntimeFeatures(), - ) - if issubclass(runner_type, PoolingRunner): - kwargs["model_caller"] = model_caller or Mock(return_value={"logits": torch.empty(0)}) - with patch.object(InputBuffers, "allocate", return_value=buffers): - return runner_type(model, config, **kwargs) - - -def _config() -> NoKVCacheRunnerConfig: - return NoKVCacheRunnerConfig( - max_batch_size=4, - max_num_tokens=16, - max_seq_len=16, - max_beam_width=1, - without_logits=False, - attention_backend=_AttentionBackend, - enable_attention_dp=False, - prefill_cuda_graph_backend=PrefillCudaGraphBackend.DISABLED, - prefill_cuda_graph_num_tokens=[], - mm_encoder_cache_enabled=False, - ) - - -def _model(*, is_generation: bool, is_encoder_decoder: bool = False) -> SimpleNamespace: - pretrained_config = SimpleNamespace( - num_attention_heads=8, - num_key_value_heads=2, - ) - return SimpleNamespace( - extra_attrs={}, - model_config=SimpleNamespace( - is_generation=is_generation, - is_encoder_decoder=is_encoder_decoder, - pretrained_config=pretrained_config, - sparse_attention_config=None, - enable_flash_mla=False, - ), - config=pretrained_config, - ) - - -@pytest.mark.parametrize( - ("encode_only", "mm_encoder_only", "is_generation", "is_encoder_decoder", "runner_type"), - [ - (False, True, False, False, MultimodalEncoderRunner), - (False, False, False, False, PoolingRunner), - (True, False, False, False, EncoderRunner), - (False, False, True, True, EncoderDecoderRunner), - (False, False, True, False, DecoderRunner), - ], -) -def test_resolve_runner_dispatches_startup_families( - encode_only: bool, - mm_encoder_only: bool, - is_generation: bool, - is_encoder_decoder: bool, - runner_type: type[ModelRunner], -) -> None: - args = SimpleNamespace(encode_only=encode_only, mm_encoder_only=mm_encoder_only) - - actual_type = resolve_runner_type( - _model(is_generation=is_generation, is_encoder_decoder=is_encoder_decoder), - args, - ) - - assert actual_type is runner_type - - -def test_resolve_runner_type_does_not_require_runtime_dependencies() -> None: - args = SimpleNamespace(encode_only=False, mm_encoder_only=False) - - assert resolve_runner_type(_model(is_generation=False), args) is PoolingRunner - - -def test_resolve_runner_checks_encode_only_before_pooling() -> None: - args = SimpleNamespace(encode_only=True, mm_encoder_only=False) - - assert resolve_runner_type(_model(is_generation=False), args) is EncoderRunner - - -def test_resolve_runner_checks_mm_encoder_before_non_generation() -> None: - args = SimpleNamespace(encode_only=False, mm_encoder_only=True) - - assert resolve_runner_type(_model(is_generation=False), args) is MultimodalEncoderRunner - - -@pytest.mark.parametrize( - ("runner_type", "initializer_name"), - [ - (EncoderRunner, "_initialize_encoder_runner"), - (EncoderDecoderRunner, "_initialize_encoder_decoder_runner"), - (NoKVCacheRunner, "_initialize_no_kv_cache_runner"), - (DecoderRunner, "_initialize_decoder_runner"), - ], -) -def test_model_engine_initializes_runner_by_family( - runner_type: type[ModelRunner], initializer_name: str -) -> None: - engine = object.__new__(ModelEngine) - engine._model_caller = Mock() - expected = Mock(spec=runner_type) - initializer = Mock(return_value=expected) - setattr(engine, initializer_name, initializer) - - assert engine._initialize_runner(runner_type) is expected - initializer.assert_called_once_with(runner_type) - - -@pytest.mark.parametrize( - ("cuda_graph_config", "expected_encode_config"), - [ - (None, False), - (CudaGraphConfig(), False), - (EncodeCudaGraphConfig(batch_sizes=[1], num_tokens=[16], seq_lens=[16]), True), - ], -) -def test_encoder_runner_graph_config_comes_only_from_cuda_graph_config( - cuda_graph_config: CudaGraphConfig | None, - expected_encode_config: bool, -) -> None: - """Encode-only takes its buckets from `cuda_graph_config`, never from llm_args.""" - engine = object.__new__(ModelEngine) - engine._model_caller = Mock() - engine.model = object() - engine.mapping = object() - engine.cuda_graph_config = cuda_graph_config - engine.llm_args = SimpleNamespace( - encoder_cuda_graph_config=EncodeCudaGraphConfig( - batch_sizes=[8], - num_tokens=[64], - seq_lens=[64], - ), - enable_autotuner=False, - ) - engine.batch_size = 4 - engine.max_num_tokens = 16 - engine.max_seq_len = 16 - engine.max_beam_width = 1 - engine.without_logits = False - engine.attn_backend = _AttentionBackend - engine.attn_runtime_features = AttentionRuntimeFeatures() - engine.dist = object() - engine.moe_load_balancer = None - runner_config = object() - expected_runner = object() - runner_type = Mock(return_value=expected_runner) - - with patch( - "tensorrt_llm._torch.pyexecutor.model_engine.EncoderRunnerConfig.create", - return_value=runner_config, - ) as create_config: - actual = engine._initialize_encoder_runner(runner_type) - - assert actual is expected_runner - resolved = create_config.call_args.kwargs["graph_config"] - assert resolved is (cuda_graph_config if expected_encode_config else None) - runner_type.assert_called_once_with( - engine.model, - runner_config, - mapping=engine.mapping, - dist=engine.dist, - moe_load_balancer=engine.moe_load_balancer, - model_caller=engine._model_caller, - attention_runtime_features=engine.attn_runtime_features, - ) - - -def test_model_engine_rejects_unregistered_runner_family() -> None: - class UnregisteredRunner: - pass - - engine = object.__new__(ModelEngine) - engine._model_caller = Mock() - - with pytest.raises(TypeError, match="No runner initializer registered"): - engine._initialize_runner(UnregisteredRunner) - - -def _model_engine_with_runner( - runner: Mock | None, - *, - kv_cache_manager: object | None, -) -> tuple[ModelEngine, Mock]: - engine = object.__new__(ModelEngine) - engine._model_caller = Mock() - engine.model = SimpleNamespace(extra_attrs={}) - if runner is None: - runner = object.__new__(DecoderRunner) - runner._model = engine.model - runner._trtllm_gen_jit_warmup = False - engine._runner = runner - engine._warmup_depth = 0 - engine.moe_load_balancer = None - resource_manager = Mock() - resource_manager.get_resource_manager.return_value = kv_cache_manager - return engine, resource_manager - - -def test_model_engine_forward_delegates_to_resolved_runner() -> None: - runner = Mock(spec=NoKVCacheRunner) - expected_outputs = {"logits": object()} - runner.forward.return_value = expected_outputs - engine, resource_manager = _model_engine_with_runner( - runner, - kv_cache_manager=None, - ) - batch = ScheduledRequests() - - actual_outputs = engine.forward(batch, resource_manager) - - assert actual_outputs is expected_outputs - runner.forward.assert_called_once_with( - ScheduledInputs(batch=batch), - resource_manager=resource_manager, - is_dummy=False, - ) - - -def test_no_kv_cache_runner_rejects_kv_manager_before_preparation() -> None: - runner = _make_runner(PoolingRunner, _model(is_generation=False)) - runner._prepare_inputs = Mock() - resource_manager = Mock() - resource_manager.get_resource_manager.return_value = object() - - with pytest.raises( - AssertionError, - match="no-KV-cache runner was initialized, but a KV cache manager was allocated", - ): - runner.forward( - ScheduledInputs(batch=ScheduledRequests()), resource_manager=resource_manager - ) - - runner._prepare_inputs.assert_not_called() - - -def test_model_engine_forward_encoder_delegates_scheduled_encoder_batch() -> None: - runner = Mock(spec=EncoderDecoderRunner) - hidden_states = object() - runner.forward_encoder.return_value = { - "encoder_hidden_states": hidden_states, - "encoder_seq_lens": [2, 3], - } - engine = object.__new__(ModelEngine) - engine._model_caller = Mock() - engine.model = SimpleNamespace( - model_config=SimpleNamespace(is_encoder_decoder=True), - ) - engine._runner = runner - engine._warmup_depth = 0 - requests = [object(), object()] - resource_manager = object() - - outputs = engine.forward_encoder(requests, resource_manager) - - assert outputs == (hidden_states, [2, 3]) - inputs = runner.forward_encoder.call_args.args[0] - assert inputs.batch.encoder_requests == requests - runner.forward_encoder.assert_called_once_with( - inputs, - resource_manager=resource_manager, - ) - - -def test_model_engine_releases_runner_owned_graphs() -> None: - engine = object.__new__(ModelEngine) - engine._model_caller = Mock() - engine._runner = Mock(spec=EncoderRunner) - engine._torch_compile_backend = None - - engine.release_executor_resources(release_model_owned_graphs=False) - - engine._runner.release_graphs.assert_called_once_with(release_nccl_window_owners=True) - - -def test_prepared_inputs_is_frozen_and_preserves_kwargs_identity() -> None: - kwargs = {"input_ids": torch.tensor([1])} - prepared = PreparedInputs(kwargs) - - assert prepared.kwargs is kwargs - with pytest.raises(FrozenInstanceError): - prepared.kwargs = {} - - -def test_no_kv_cache_config_extends_common_runner_config() -> None: - assert isinstance(_config(), RunnerConfig) - - -def test_no_kv_cache_runner_exposes_only_forward_step_as_abstract() -> None: - assert NoKVCacheRunner.__abstractmethods__ == frozenset({"_forward_step"}) - - -@pytest.mark.parametrize("runner_type", [PoolingRunner, MultimodalEncoderRunner]) -def test_concrete_runners_implement_no_kv_cache_forward_step( - runner_type: type[NoKVCacheRunner], -) -> None: - assert issubclass(runner_type, NoKVCacheRunner) - assert runner_type._forward_step is not NoKVCacheRunner._forward_step - - -def test_no_kv_cache_runner_owns_and_reuses_attention_metadata() -> None: - runner = _make_runner(PoolingRunner, _model(is_generation=False)) - - first = runner._setup_attn_metadata() - second = runner._setup_attn_metadata() - - assert first is second - assert first.kv_cache_manager is None - assert first.max_num_requests == 4 - assert first.max_num_sequences == 4 - assert first.max_num_heads_per_kv == 4 - assert first.block_ids_per_seq is None - assert first.kv_block_ids_per_seq is None - - -def test_pooling_runner_owns_forward_output_processing( - monkeypatch: pytest.MonkeyPatch, -) -> None: - logits = torch.arange(12).reshape(3, 4) - model_forward = Mock(return_value=logits) - runner = _make_runner( - PoolingRunner, - _model(is_generation=False), - model_caller=model_forward, - ) - attn_metadata = SimpleNamespace(on_update_kv_lens=Mock()) - prepared = PreparedInputs( - {"attn_metadata": attn_metadata, "input_ids": torch.tensor([1])}, - ) - monkeypatch.setattr(runner, "_prepare_inputs", Mock(return_value=prepared)) - - outputs = runner.forward( - ScheduledInputs(batch=ScheduledRequests()), - resource_manager=SimpleNamespace(get_resource_manager=Mock(return_value=None)), - ) - - torch.testing.assert_close(outputs["logits"], logits) - attn_metadata.on_update_kv_lens.assert_called_once_with() - model_forward.assert_called_once_with( - attn_metadata=attn_metadata, - input_ids=prepared.kwargs["input_ids"], - return_context_logits=False, - ) - - -def test_pooling_runner_returns_raw_model_outputs_when_logits_are_disabled() -> None: - model_outputs = {"hidden_states": object()} - model_forward = Mock(return_value=model_outputs) - runner = _make_runner( - PoolingRunner, - _model(is_generation=False), - replace(_config(), without_logits=True), - model_caller=model_forward, - ) - - outputs = runner._forward_step({}, SimpleNamespace()) - - assert outputs is model_outputs - - -@pytest.mark.parametrize("runner_type", [PoolingRunner, MultimodalEncoderRunner]) -def test_no_kv_cache_runner_default_lifecycle_with_no_kv_resources( - runner_type: type[PoolingRunner] | type[MultimodalEncoderRunner], -) -> None: - runner = _make_runner(runner_type, _model(is_generation=False)) - resource_manager = SimpleNamespace(get_resource_manager=Mock(return_value=None)) - - assert runner.warmup(resource_manager) is None - assert runner.release_graphs() is None - assert runner.wait_for_input_copy() is None - - -def test_mm_encoder_runner_forward_step_returns_empty_result_without_multimodal_params() -> None: - runner = _make_runner(MultimodalEncoderRunner, SimpleNamespace()) - - result = runner._forward_step({}, SimpleNamespace(context_requests=[])) - - assert result == { - "mm_embeddings": [], - "mm_embedding_request_indices": [], - "mm_embedding_lengths": [], - } - - -def test_mm_encoder_runner_forward_step_rejects_payload_request_count_mismatch() -> None: - request = SimpleNamespace(py_multimodal_data={"image": object()}) - runner = _make_runner(MultimodalEncoderRunner, SimpleNamespace()) - - with pytest.raises(ValueError, match="one multimodal payload per context"): - runner._forward_step( - {"multimodal_params": [SimpleNamespace(), SimpleNamespace()]}, - SimpleNamespace(context_requests=[request]), - ) - - -def test_mm_encoder_runner_forward_step_skips_metadata_only_and_missing_length_requests() -> None: - requests = [ - SimpleNamespace(py_multimodal_data={"mrope_config": {}}), - SimpleNamespace(py_multimodal_data={"image": object()}), - SimpleNamespace( - py_multimodal_data={ - "image": object(), - MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( - [("image", 0)], output_embedding_lengths=[2] - ), - }, - multimodal_lengths=[2], - ), - ] - params = [ - SimpleNamespace(multimodal_data={}), - SimpleNamespace(multimodal_data={}), - SimpleNamespace(multimodal_data={}), - ] - model = SimpleNamespace(forward=Mock(return_value=[torch.arange(4).reshape(2, 2)])) - runner = _make_runner(MultimodalEncoderRunner, model) - - result = runner._forward_step( - {"multimodal_params": params}, - SimpleNamespace(context_requests=requests), - ) - - model.forward.assert_called_once_with([params[2]]) - assert result["logits"] is None - assert result["mm_embedding_request_indices"] == [2] - assert result["mm_embedding_lengths"] == [[2]] - torch.testing.assert_close(result["mm_embeddings"][0], torch.arange(4).reshape(2, 2)) - - -def test_mm_encoder_runner_splits_embeddings_and_returns_mrope_metadata() -> None: - first_param = SimpleNamespace( - multimodal_data={ - "image": object(), - "mrope_config": { - "mrope_position_ids": "first-ids", - "mrope_position_deltas": "first-deltas", - }, - } - ) - second_param = SimpleNamespace( - multimodal_data={ - "image": object(), - "mrope_config": { - "mrope_position_ids": "second-ids", - "mrope_position_deltas": "second-deltas", - }, - } - ) - requests = [ - SimpleNamespace( - py_multimodal_data={ - "image": object(), - MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( - [("image", 0)], output_embedding_lengths=[1] - ), - }, - multimodal_lengths=[1], - ), - SimpleNamespace( - py_multimodal_data={ - "image": object(), - MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( - [("image", 0)], output_embedding_lengths=[2] - ), - }, - multimodal_lengths=[2], - ), - ] - embeddings = torch.arange(6).reshape(3, 2) - model = SimpleNamespace(forward=Mock(return_value=[embeddings])) - runner = _make_runner(MultimodalEncoderRunner, model) - - result = runner._forward_step( - {"multimodal_params": [first_param, second_param]}, - SimpleNamespace(context_requests=requests), - ) - - model.forward.assert_called_once_with([first_param, second_param]) - assert result["logits"] is None - assert result["mm_embedding_request_indices"] == [0, 1] - assert result["mm_embedding_lengths"] == [[1], [2]] - torch.testing.assert_close(result["mm_embeddings"][0], embeddings[:1]) - torch.testing.assert_close(result["mm_embeddings"][1], embeddings[1:]) - assert result["mrope_position_ids"] == ["first-ids", "second-ids"] - assert result["mrope_position_deltas"] == ["first-deltas", "second-deltas"] - - -@pytest.mark.parametrize("effective_length", [None, 0, 3]) -@pytest.mark.parametrize("is_dummy", [False, True]) -def test_engine_passes_call_state_and_returns_length_update(effective_length, is_dummy): - runner = Mock(spec=NoKVCacheRunner) - outputs = {"logits": object()} - if effective_length is not None: - outputs["runtime_draft_len"] = effective_length - runner.forward.return_value = outputs - engine, resources = _model_engine_with_runner(runner, kv_cache_manager=None) - engine._warmup_depth = int(is_dummy) - batch = ScheduledRequests() - batch.context_requests_last_chunk = [SimpleNamespace(py_return_context_logits=True)] - - assert engine.forward(batch, resources, enable_spec_decode=True, runtime_draft_len=2) is outputs - - inputs = runner.forward.call_args.args[0] - assert inputs.batch is batch - assert inputs.gather_context_logits - assert inputs.enable_spec_decode - assert inputs.runtime_draft_len == 2 - assert runner.forward.call_args.kwargs["is_dummy"] is is_dummy - assert outputs.get("runtime_draft_len") == effective_length - - -@pytest.mark.parametrize("is_dummy", [False, True]) -def test_decoder_runner_passes_call_state_without_engine_writes(is_dummy): - engine, resources = _model_engine_with_runner(None, kv_cache_manager=object()) - # The dummy flag comes from the caller, not from the warmup scope. - engine._warmup_depth = int(not is_dummy) - inputs = ScheduledInputs( - batch=ScheduledRequests(), enable_spec_decode=True, runtime_draft_len=2 - ) - expected = {"logits": object(), "runtime_draft_len": 4} - engine._runner._forward_decoder = Mock(return_value=expected) - - outputs = engine._forward_scheduled(inputs, resource_manager=resources, is_dummy=is_dummy) - - assert outputs is expected - engine._runner._forward_decoder.assert_called_once_with(inputs, resources, is_dummy=is_dummy) - assert (inputs.enable_spec_decode, inputs.runtime_draft_len) == (True, 2) - - -def test_forward_warmup_runs_local_call_state_in_forward_context(): - engine, resources = _model_engine_with_runner(None, kv_cache_manager=object()) - batch = ScheduledRequests() - - def decoder_forward(inputs, resource_manager, *, is_dummy): - assert torch.is_inference_mode_enabled() - assert get_model_extra_attrs() is engine.model.extra_attrs - return {"logits": object(), "runtime_draft_len": 3} - - engine._runner._forward_decoder = Mock(side_effect=decoder_forward) - - outputs = engine._runner._forward_warmup( - batch, resources, enable_spec_decode=True, runtime_draft_len=1 - ) - - inputs = engine._runner._forward_decoder.call_args.args[0] - assert inputs.batch is batch - assert inputs.new_tensors_device is None - assert inputs.cache_indirection_buffer is None - assert (inputs.enable_spec_decode, inputs.runtime_draft_len) == (True, 1) - assert engine._runner._forward_decoder.call_args.kwargs == {"is_dummy": True} - assert outputs["runtime_draft_len"] == 3 - assert get_model_extra_attrs() is None - - -def test_decoder_runner_forward_enters_model_context(): - engine, resources = _model_engine_with_runner(None, kv_cache_manager=object()) - inputs = ScheduledInputs(batch=ScheduledRequests()) - - def decoder_forward(inputs, resource_manager, *, is_dummy): - assert torch.is_inference_mode_enabled() - assert get_model_extra_attrs() is engine.model.extra_attrs - return {"logits": object()} - - engine._runner._forward_decoder = Mock(side_effect=decoder_forward) - - engine._runner.forward(inputs, resource_manager=resources) - - engine._runner._forward_decoder.assert_called_once_with(inputs, resources, is_dummy=False) - assert get_model_extra_attrs() is None - - -def test_no_kv_forward_enters_model_context(): - model = _model(is_generation=False) - - def model_forward(**kwargs): - assert torch.is_inference_mode_enabled() - assert get_model_extra_attrs() is model.extra_attrs - return {"hidden_states": object()} - - caller = Mock(side_effect=model_forward) - runner = _make_runner(PoolingRunner, model, model_caller=caller) - runner._prepare_inputs = Mock(return_value=PreparedInputs({})) - resources = SimpleNamespace(get_resource_manager=Mock(return_value=None)) - - runner.forward(ScheduledInputs(batch=ScheduledRequests()), resource_manager=resources) - - caller.assert_called_once() - assert get_model_extra_attrs() is None - - -def test_no_kv_forward_preserves_model_output_dictionary_without_length_update(): - outputs = {"hidden_states": object()} - caller = Mock(return_value=outputs) - runner = _make_runner(PoolingRunner, _model(is_generation=False), model_caller=caller) - runner._prepare_inputs = Mock(return_value=PreparedInputs({})) - inputs = ScheduledInputs(batch=ScheduledRequests(), runtime_draft_len=3) - resources = SimpleNamespace(get_resource_manager=Mock(return_value=None)) - - actual = runner.forward(inputs, resource_manager=resources) - - assert actual is outputs - assert "runtime_draft_len" not in outputs - assert inputs.runtime_draft_len == 3 - - -@pytest.mark.parametrize("raw_output", [torch.tensor([1.0]), None], ids=["tensor", "none"]) -@pytest.mark.parametrize("is_dummy", [False, True]) -def test_engine_forward_preserves_raw_decoder_outputs(raw_output, is_dummy): - engine, resources = _model_engine_with_runner(None, kv_cache_manager=object()) - engine._warmup_depth = int(is_dummy) - engine._runner._forward_decoder = Mock(return_value=raw_output) - - assert engine.forward(ScheduledRequests(), resources, runtime_draft_len=2) is raw_output - assert engine._runner._forward_decoder.call_args.args[0].runtime_draft_len == 2 - assert engine._runner._forward_decoder.call_args.kwargs == {"is_dummy": is_dummy} - - -def test_model_caller_uses_current_forward_and_restores_outer_attribute_context(): - metadata = _AttentionMetadata() - spec_metadata = object() - attrs = {} - outer_attrs = {"outer": True} - model = SimpleNamespace(model_config=SimpleNamespace(extra_attrs={"model_flag": True})) - - def first_forward(**kwargs): - assert get_model_extra_attrs() is attrs - assert attrs["attention_metadata"]() is metadata - assert attrs["spec_metadata"] is spec_metadata - assert attrs["model_flag"] is True - return "first" - - model.forward = Mock(side_effect=first_forward) - caller = ModelCaller(model) - with model_extra_attrs(outer_attrs): - with model_extra_attrs(attrs): - assert caller(attn_metadata=metadata, spec_metadata=spec_metadata) == "first" - model.forward = Mock(return_value="replacement") - assert caller(attn_metadata=metadata) == "replacement" - assert attrs["spec_metadata"] is None - model.forward.side_effect = RuntimeError("model failure") - with pytest.raises(RuntimeError, match="model failure"): - caller(attn_metadata=metadata) - assert get_model_extra_attrs() is outer_attrs - - -def test_model_caller_borrows_live_compile_streams_and_events(): - streams = Backend.Streams() - backend = SimpleNamespace(events=Backend.Events()) - metadata = _AttentionMetadata() - current_stream = object() - attrs = {} - - def forward(**kwargs): - assert get_model_extra_attrs() is attrs - assert attrs["aux_streams"]() is streams - assert attrs["events"]() is backend.events - assert attrs["global_stream"] is current_stream - return len(attrs["aux_streams"]()), len(attrs["events"]()) - - model = SimpleNamespace(model_config=SimpleNamespace(extra_attrs={}), forward=forward) - caller = ModelCaller(model, compile_backend=backend, aux_streams=streams) - with patch("torch.cuda.current_stream", return_value=current_stream), model_extra_attrs(attrs): - assert caller(attn_metadata=metadata) == (0, 0) - streams.append(object()) - backend.events = Backend.Events([object(), object()]) - assert caller(attn_metadata=metadata) == (1, 2) - - -def test_model_caller_requires_streams_with_compile_backend(): - with pytest.raises(ValueError, match="requires its auxiliary stream container"): - ModelCaller(SimpleNamespace(), compile_backend=SimpleNamespace(events=Backend.Events())) - - -def test_mm_encoder_handoff_preserves_per_item_metadata_in_sparse_batch() -> None: - requests = [ - SimpleNamespace( - py_multimodal_data={ - "image": object(), - MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( - [("image", 0)], output_embedding_lengths=[1] - ), - }, - multimodal_lengths=[1], - ), - SimpleNamespace(py_multimodal_data=None), - SimpleNamespace( - py_multimodal_data={ - "video": object(), - MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( - [("video", 0), ("video", 1)], output_embedding_lengths=[3, 2] - ), - }, - multimodal_lengths=[7, 6], - ), - ] - for request in requests: - request.py_result = PyResult(prompt_len=1, max_new_tokens=1) - request.set_finished_reason = Mock() - params = [ - SimpleNamespace(multimodal_data=request.py_multimodal_data) - for request in requests - if request.py_multimodal_data is not None - ] - scheduled_requests = SimpleNamespace( - context_requests=requests, generation_requests=[], num_context_requests=len(requests) - ) - embeddings = torch.arange(12, dtype=torch.float32).reshape(6, 2) - - def encode(batch_params: list[SimpleNamespace]) -> list[torch.Tensor]: - batch_params[1].multimodal_data["multimodal_embedding_metadata"] = [ - {"retained_token_counts": [2, 1]}, - {"retained_token_counts": [2, 0]}, - ] - return [embeddings] - - runner = _make_runner(MultimodalEncoderRunner, SimpleNamespace(forward=encode)) - result = runner._forward_step( - {"multimodal_params": params}, - scheduled_requests, - ) - sampler = EarlyStopWithMMResult() - state = sampler.sample_async(scheduled_requests, result, []) - sampler.update_requests(state) - - assert requests[1].py_result.mm_embedding_handles is None - image_handles = json.loads(json.dumps(requests[0].py_result.mm_embedding_handles)) - assert len(image_handles) == 1 - assert "metadata" not in image_handles[0] - torch.testing.assert_close( - SharedTensorContainer.from_dict(image_handles[0]).get_local_view(), embeddings[:1] - ) - video_handles = json.loads(json.dumps(requests[2].py_result.mm_embedding_handles)) - assert [handle["metadata"]["retained_token_counts"] for handle in video_handles] == [ - [2, 1], - [2, 0], - ] - restored = [ - SharedTensorContainer.from_dict(handle).get_local_view() for handle in video_handles - ] - assert [len(embedding) for embedding in restored] == [3, 2] - torch.testing.assert_close(torch.cat(restored), embeddings[1:]) diff --git a/tests/unittest/_torch/executor/engine_test_utils.py b/tests/unittest/_torch/executor/engine_test_utils.py new file mode 100644 index 000000000000..e67ed50d51f0 --- /dev/null +++ b/tests/unittest/_torch/executor/engine_test_utils.py @@ -0,0 +1,311 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + + +from __future__ import annotations + +import contextlib +from dataclasses import dataclass +from types import SimpleNamespace +from typing import TYPE_CHECKING, Any +from unittest.mock import Mock, patch + +import torch + +import tensorrt_llm +import tensorrt_llm._torch.pyexecutor.engine.runners.decoder.warmup as decoder_warmup_module +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import DecoderRunner +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.warmup import WarmupAgreement +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest +from tensorrt_llm._torch.pyexecutor.model_engine import ModelEngine +from tensorrt_llm.llmapi import CudaGraphConfig, KvCacheConfig, SamplingParams +from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig, TorchLlmArgs +from tensorrt_llm.mapping import Mapping + +if TYPE_CHECKING: + from collections.abc import Iterator + + +@dataclass +class Config: + torch_dtype: torch.dtype + num_key_value_heads: int = 16 + num_attention_heads: int = 16 + hidden_size: int = 256 + architectures: list[str] | None = None + rope_scaling: dict[str, object] | None = None + + @property + def head_dim(self) -> int: + return self.hidden_size // self.num_attention_heads + + +class DummyModel(torch.nn.Module): + def __init__(self, dtype: torch.dtype, rope_scaling: dict[str, object] | None = None) -> None: + super().__init__() + self.model_config = ModelConfig( + pretrained_config=Config(torch_dtype=dtype, rope_scaling=rope_scaling) + ) + self.recorded_position_ids: torch.Tensor | None = None + self.recorded_kwargs: list[dict[str, Any]] = [] + + def infer_max_seq_len(self) -> int: + return 2048 + + @property + def config(self) -> Config: + return self.model_config.pretrained_config + + def forward(self, *args: object, **kwargs: Any) -> dict[str, torch.Tensor]: + input_ids = kwargs["input_ids"] + self.recorded_position_ids = kwargs["position_ids"] + # Input buffers are reused across steps, so keep copies. + if not torch.cuda.is_current_stream_capturing(): + self.recorded_kwargs.append(_clone_tensors(kwargs)) + batch_size = input_ids.size(0) + return {"logits": torch.randn((batch_size, 10), device="cuda")} + + +class DummyModelEngine(ModelEngine): + def __init__( + self, + llm_args: TorchLlmArgs, + dtype: torch.dtype, + spec_config: DecodingBaseConfig | None = None, + is_generation: bool = True, + model: DummyModel | None = None, + ) -> None: + self.dtype = dtype + mapping = Mapping( + world_size=tensorrt_llm.mpi_world_size(), + tp_size=tensorrt_llm.mpi_world_size(), + rank=tensorrt_llm.mpi_rank(), + ) + if model is None: + model = DummyModel(self.dtype) + model.model_config.is_generation = is_generation + super().__init__( + model_path="dummy", + mapping=mapping, + model=model, + llm_args=llm_args, + spec_config=spec_config, + ) + + +def _create_request(num_tokens: int, req_id: int) -> LlmRequest: + sampling_params = SamplingParams() + kwargs = { + "request_id": req_id, + "max_new_tokens": 1, + "input_tokens": [0] * num_tokens, + "sampling_config": tensorrt_llm.bindings.SamplingConfig( + sampling_params._get_sampling_config() + ), + "is_streaming": False, + } + result = LlmRequest(**kwargs) + result.paged_kv_block_ids = [] + return result + + +def _create_request_with_tokens(tokens: list[int], req_id: int) -> LlmRequest: + sampling_params = SamplingParams() + request = LlmRequest( + request_id=req_id, + max_new_tokens=1, + input_tokens=tokens, + sampling_config=tensorrt_llm.bindings.SamplingConfig( + sampling_params._get_sampling_config() + ), + is_streaming=False, + ) + request.paged_kv_block_ids = [] + return request + + +def _make_request_stub(req_id: int, prompt_len: int = 4) -> SimpleNamespace: + return SimpleNamespace( + py_request_id=req_id, + context_chunk_size=1, + context_remaining_length=1, + context_current_position=prompt_len - 1, + py_prompt_len=prompt_len, + py_beam_width=1, + py_draft_tokens=[], + is_context_only_request=False, + is_generation_only_request=False, + py_disaggregated_params=None, + py_multimodal_data=None, + py_mm_encoder_event=None, + py_mrope_position_delta=None, + py_return_context_logits=False, + py_batch_idx=None, + lora_task_id=None, + is_dummy=False, + max_beam_num_tokens=prompt_len, + state="context", + py_llm_request_type="context_and_generation", + ) + + +@contextlib.contextmanager +def create_model_engine_and_kvcache( + llm_args: TorchLlmArgs | None = None, + execution_stream: torch.cuda.Stream | None = None, + spec_config: DecodingBaseConfig | None = None, + model: DummyModel | None = None, +) -> Iterator[tuple[DummyModelEngine, KVCacheManagerV2]]: + if llm_args is None: + llm_args = TorchLlmArgs( + model="dummy", + max_batch_size=13, + max_num_tokens=258, + cuda_graph_config=CudaGraphConfig( + enable_padding=True, batch_sizes=[1, 2, 4, 8, 16, 32, 64, 128] + ), + ) + assert all(size not in llm_args.cuda_graph_config.batch_sizes for size in (5, 13)) + assert all(size in llm_args.cuda_graph_config.batch_sizes for size in (8, 16)) + engine = DummyModelEngine(llm_args, torch.half, spec_config, model=model) + with kv_cache_manager(model=engine.model, execution_stream=execution_stream) as manager: + try: + yield engine, manager + finally: + engine.release_executor_resources(release_model_owned_graphs=True) + + +@contextlib.contextmanager +def kv_cache_manager( + *, + model: DummyModel, + max_tokens: int = 258, + max_batch_size: int = 13, + mapping: Mapping | None = None, + execution_stream: torch.cuda.Stream | None = None, + config: KvCacheConfig | None = None, + max_num_tokens: int | None = None, +) -> Iterator[KVCacheManagerV2]: + options = {} if max_num_tokens is None else {"max_num_tokens": max_num_tokens} + manager = KVCacheManagerV2( + config or KvCacheConfig(max_tokens=max_tokens), + tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF, + num_layers=1, + num_kv_heads=model.config.num_key_value_heads, + head_dim=model.config.head_dim, + tokens_per_block=1, + max_seq_len=max_tokens, + max_batch_size=max_batch_size, + mapping=mapping or Mapping(world_size=1, tp_size=1, rank=0), + dtype=tensorrt_llm.bindings.DataType.HALF, + execution_stream=execution_stream, + **options, + ) + try: + yield manager + finally: + manager.shutdown() + + +def _clone_tensors(values: dict[str, Any]) -> dict[str, Any]: + return { + name: value.clone() if isinstance(value, torch.Tensor) else value + for name, value in values.items() + } + + +@contextlib.contextmanager +def _capture_tllm_logs() -> Iterator[list[str]]: + """Capture logger.info calls emitted from the decoder runner module. + + tensorrt_llm.logger.logger is a custom Singleton (not a stdlib + logging.Logger) and does not route through stdlib logging by default, + so a logging.Handler attached to logging.getLogger("tensorrt_llm") + sees nothing. Patch the logger.info bound on the decoder runner module + directly so we observe exactly the messages warmup() emits. + """ + records = [] + + def _record(*msg: object) -> None: + records.append(" ".join(str(m) for m in msg)) + + with patch.object(decoder_warmup_module.logger, "info", side_effect=_record): + yield records + + +def _agreement( + *, + world_size: int = 1, + dwdp_size: int = 0, + tp_size: int = 1, + tp_peer_values: tuple[int, ...] = (), + peer_values: tuple[int, ...] | None = None, + with_dist: bool = True, +) -> WarmupAgreement: + """Build the warmup agreement over a stand-in communicator and mapping. + + ``dist`` is ``Optional`` on the real engine -- stub engines such as + ``DummyModelEngine`` construct without a communicator -- so ``with_dist`` + covers that state too. + + ``tp_allgather`` returns this rank's value followed by ``tp_peer_values``; + ``allgather`` does the same over ``peer_values`` for the forward group. + ``peer_values`` defaults to peers that mirror this rank, i.e. a world that + agrees with whatever this rank reports. + """ + + def _allgather(value: int) -> list[int]: + peers = [value] * (world_size - 1) if peer_values is None else list(peer_values) + return [value, *peers] + + dist = ( + SimpleNamespace( + world_size=world_size, + rank=3, + tp_allgather=lambda value: [value, *tp_peer_values], + allgather=_allgather, + ) + if with_dist + else None + ) + mapping = SimpleNamespace( + dwdp_enabled=dwdp_size > 1, + has_cp_helix=lambda: False, + tp_size=tp_size, + ) + return WarmupAgreement(dist, mapping) + + +class _AttentionMetadata: + def __init__(self, **kwargs: object) -> None: + self.__dict__.update(kwargs) + self.block_ids_per_seq = object() + self.kv_block_ids_per_seq = object() + self.on_update_kv_lens = Mock() + + +class _AttentionBackend: + Metadata = _AttentionMetadata + + +def _model_engine_with_runner( + runner: Mock | None, + *, + kv_cache_manager: object | None, +) -> tuple[ModelEngine, Mock]: + engine = object.__new__(ModelEngine) + engine._cleanup_done = True + engine._model_caller = Mock() + engine.model = SimpleNamespace(extra_attrs={}) + if runner is None: + runner = object.__new__(DecoderRunner) + runner._model = engine.model + runner._trtllm_gen_jit_warmup = False + engine._runner = runner + engine._warmup_depth = 0 + engine.moe_load_balancer = None + resource_manager = Mock() + resource_manager.get_resource_manager.return_value = kv_cache_manager + return engine, resource_manager diff --git a/tests/unittest/_torch/executor/test_cuda_graph_runner.py b/tests/unittest/_torch/executor/test_cuda_graph_runner.py index 42d5e5da7de5..1fa9c38800ad 100644 --- a/tests/unittest/_torch/executor/test_cuda_graph_runner.py +++ b/tests/unittest/_torch/executor/test_cuda_graph_runner.py @@ -1,26 +1,314 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + import contextlib -from unittest.mock import MagicMock, Mock +import unittest +from types import SimpleNamespace +from typing import TYPE_CHECKING +from unittest.mock import MagicMock, Mock, patch import pytest import torch +from _torch.executor.engine_test_utils import _make_request_stub import tensorrt_llm._torch.pyexecutor.cuda_graph_runner as cuda_graph_runner_module from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import ( + CUDAGraphRunner, EncoderCUDAGraphRunner, EncoderCUDAGraphRunnerConfig, + KeyType, + SampleType, + _restore_spec_decode_capture_state, + _save_spec_decode_capture_state, ) +from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests +from tensorrt_llm.llmapi.llm_args import SeqLenAwareSparseAttentionConfig + +if TYPE_CHECKING: + from collections.abc import Iterator pytestmark = pytest.mark.cpu_only -def test_encoder_capture_discards_warmup_and_capture_outputs(monkeypatch) -> None: +def _feature_encoder_runner( + batch_sizes: list[int], fixed_seq_len: int = 1_500 +) -> EncoderCUDAGraphRunner: + """Build a feature-mode backend without enabling CUDA capture.""" + config = EncoderCUDAGraphRunnerConfig( + use_cuda_graph=False, + cuda_graph_padding_enabled=True, + cuda_graph_batch_sizes=batch_sizes, + cuda_graph_num_tokens=[], + cuda_graph_seq_lens=[], + max_cuda_graph_batch_size=max(batch_sizes), + max_cuda_graph_num_tokens=max(batch_sizes) * fixed_seq_len, + max_num_tokens=max(batch_sizes) * fixed_seq_len, + max_seq_len=fixed_seq_len, + is_encoder_decoder=True, + use_fixed_sequence_slots=False, + feature_shape=(480_000,), + feature_dtype=torch.float32, + fixed_seq_len=fixed_seq_len, + ) + return EncoderCUDAGraphRunner(config) + + +class GraphOwnerTeardownTestCase(unittest.TestCase): + def test_cuda_graph_runner_destructor_does_not_release_window_owner(self) -> None: + runner = Mock() + + CUDAGraphRunner.__del__(runner) + + runner.clear.assert_called_once_with(release_nccl_window_owners=False) + + +class GraphKeyTestCase(unittest.TestCase): + def test_graph_key_forwards_promoted_context_ids(self) -> None: + runner = Mock() + runner.config = SimpleNamespace() + runner._resolve_sample_type.return_value = SampleType.FULL + runner._get_seq_len_mode.return_value = True + runner._ragged_verify_bucket.return_value = None + request = _make_request_stub(7) + batch = ScheduledRequests() + batch.generation_requests = [request] + promoted_ids = frozenset({request.py_request_id}) + + key = CUDAGraphRunner.get_graph_key( + runner, + batch, + new_tensors_device=None, + promoted_context_request_ids=promoted_ids, + ) + + runner._get_seq_len_mode.assert_called_once_with(batch, None, promoted_ids) + self.assertEqual( + key, KeyType(batch_size=1, draft_len=0, short_seq_len_mode=True, num_encoder_tokens=0) + ) + + def test_graph_key_aggregates_encoder_tokens(self) -> None: + runner = Mock() + runner.config = SimpleNamespace() + runner._resolve_sample_type.return_value = SampleType.FULL + runner.max_beam_width = 1 + runner._get_seq_len_mode.return_value = False + runner._ragged_verify_bucket.return_value = None + context = _make_request_stub(1) + context.encoder_output_len = 7 + context.py_skip_cross_kv_projection = False + skipped_context = _make_request_stub(2) + skipped_context.encoder_output_len = 11 + skipped_context.py_skip_cross_kv_projection = True + batch = ScheduledRequests() + batch.context_requests_last_chunk = [context, skipped_context] + batch.generation_requests = [_make_request_stub(3)] + + key = CUDAGraphRunner.get_graph_key(runner, batch) + + assert key is not None + self.assertEqual(key.num_contexts, 2) + self.assertEqual(key.context_query_len, 1) + self.assertEqual(key.num_encoder_tokens, 7) + self.assertEqual(CUDAGraphRunner._get_num_tokens_for_key(runner, key), 3) + + def test_graph_key_rejects_nonuniform_context_query_lengths(self) -> None: + runner = Mock() + runner.config = SimpleNamespace() + runner._resolve_sample_type.return_value = SampleType.FULL + runner._get_seq_len_mode.return_value = False + runner._ragged_verify_bucket.return_value = None + first_context = _make_request_stub(1) + first_context.encoder_output_len = 7 + first_context.py_skip_cross_kv_projection = False + second_context = _make_request_stub(2) + second_context.context_chunk_size = 2 + second_context.encoder_output_len = 11 + second_context.py_skip_cross_kv_projection = False + batch = ScheduledRequests() + batch.context_requests_last_chunk = [first_context, second_context] + batch.generation_requests = [_make_request_stub(3)] + + key = CUDAGraphRunner.get_graph_key(runner, batch) + + self.assertIsNone(key) + + def test_graph_key_rounds_encoder_tokens_up_to_captured_extent(self) -> None: + key = KeyType( + batch_size=2, draft_len=0, num_contexts=1, context_query_len=1, num_encoder_tokens=7 + ) + smaller_key = key._replace(num_encoder_tokens=6) + compatible_key = key._replace(num_encoder_tokens=8) + larger_key = key._replace(num_encoder_tokens=16) + runner = SimpleNamespace( + padding_enabled=True, + _capture_allowed=False, + graph_metadata={}, + graph_outputs={ + smaller_key: object(), + compatible_key: object(), + larger_key: object(), + }, + ) + + actual_key = CUDAGraphRunner._get_compatible_mixed_encoder_decoder_key(runner, key) + + self.assertEqual(actual_key, compatible_key) + + def test_graph_key_includes_peft_cache_dtype(self) -> None: + runner = Mock() + runner.config = SimpleNamespace() + runner._resolve_sample_type.return_value = SampleType.FULL + runner._get_seq_len_mode.return_value = False + runner._ragged_verify_bucket.return_value = None + request = _make_request_stub(7) + batch = ScheduledRequests() + batch.generation_requests = [request] + + model_dtype_key = CUDAGraphRunner.get_graph_key( + runner, batch, peft_cache_data_type=torch.bfloat16 + ) + fp8_key = CUDAGraphRunner.get_graph_key( + runner, batch, peft_cache_data_type=torch.float8_e4m3fn + ) + + self.assertNotEqual(model_dtype_key, fp8_key) + self.assertEqual( + model_dtype_key._replace(peft_cache_data_type=None), + fp8_key._replace(peft_cache_data_type=None), + ) + + def test_graph_dtype_change_falls_back_to_eager(self) -> None: + runner = Mock() + runner.enabled = True + runner._agree_graph_shape = CUDAGraphRunner._agree_graph_shape.__get__(runner) + runner._dspark_confidence_enabled = False + runner.config = SimpleNamespace( + enable_attention_dp=False, + use_mrope=False, + ) + model_dtype_key = KeyType(batch_size=1, draft_len=0, peft_cache_data_type=torch.bfloat16) + fp8_key = KeyType(batch_size=1, draft_len=0, peft_cache_data_type=torch.float8_e4m3fn) + runner.get_graph_key.return_value = fp8_key + runner.graph_metadata = {model_dtype_key: object()} + runner._capture_allowed = False + runner._is_mixed_encoder_decoder_batch.return_value = False + runner._can_run_cuda_graph_batch.return_value = True + + request = _make_request_stub(7) + batch = ScheduledRequests() + batch.generation_requests = [request] + + with patch( + "tensorrt_llm._torch.pyexecutor.cuda_graph_runner.ExpertStatistic.should_record", + return_value=False, + ): + result = CUDAGraphRunner.maybe_get_cuda_graph( + runner, + batch, + enable_spec_decode=False, + attn_metadata=object(), + peft_cache_data_type=torch.float8_e4m3fn, + ) + + self.assertEqual(result, (None, None, None)) + runner.get_graph_key.assert_called_once_with( + batch, None, None, None, frozenset(), torch.float8_e4m3fn, False, None + ) + + def test_graph_key_includes_lora_variant(self) -> None: + runner = Mock() + runner.config = SimpleNamespace() + runner._resolve_sample_type.return_value = SampleType.FULL + runner._get_seq_len_mode.return_value = False + runner._ragged_verify_bucket.return_value = None + request = _make_request_stub(7) + batch = ScheduledRequests() + batch.generation_requests = [request] + + key = CUDAGraphRunner.get_graph_key( + runner, + batch, + use_lora_graph=True, + ) + + self.assertEqual( + key, + KeyType(batch_size=1, draft_len=0, use_lora_graph=True), + ) + + def test_graph_lookup_forwards_promoted_context_ids(self) -> None: + runner = Mock() + runner.enabled = True + runner._agree_graph_shape = CUDAGraphRunner._agree_graph_shape.__get__(runner) + runner._dspark_confidence_enabled = False + runner.config = SimpleNamespace( + enable_attention_dp=False, + use_mrope=False, + ) + key = KeyType(batch_size=1, draft_len=0, short_seq_len_mode=True, num_encoder_tokens=0) + graph_attn_metadata = object() + graph_spec_metadata = object() + runner.get_graph_key.return_value = key + runner.graphs = {key: object()} + runner.graph_metadata = { + key: { + "attn_metadata": graph_attn_metadata, + "spec_metadata": graph_spec_metadata, + } + } + runner._is_mixed_encoder_decoder_batch.return_value = False + request = _make_request_stub(7) + batch = ScheduledRequests() + batch.generation_requests = [request] + promoted_ids = frozenset({request.py_request_id}) + + with patch( + "tensorrt_llm._torch.pyexecutor.cuda_graph_runner.ExpertStatistic.should_record", + return_value=False, + ): + result = CUDAGraphRunner.maybe_get_cuda_graph( + runner, + batch, + enable_spec_decode=False, + attn_metadata=object(), + promoted_context_request_ids=promoted_ids, + ) + + runner.get_graph_key.assert_called_once_with( + batch, None, None, None, promoted_ids, None, False, None + ) + self.assertEqual(result, (graph_attn_metadata, graph_spec_metadata, key)) + + +class CaptureStateTestCase(unittest.TestCase): + def test_spec_decode_capture_restores_kv_lens_between_warmups(self) -> None: + attn_metadata = Mock() + attn_metadata.num_seqs = 1 + attn_metadata.kv_lens_cuda = torch.tensor([4095], dtype=torch.int32) + + saved_kv_lens_cuda = _save_spec_decode_capture_state(attn_metadata, enable_spec_decode=True) + + # CUDA graph capture performs two eager warmup forwards. A speculative + # draft loop may advance the static attention metadata during each + # forward, but the next warmup must start from the original input. + for _ in range(2): + attn_metadata.kv_lens_cuda.add_(1) + _restore_spec_decode_capture_state(attn_metadata, saved_kv_lens_cuda) + self.assertEqual(attn_metadata.kv_lens_cuda.tolist(), [4095]) + + self.assertEqual(attn_metadata.on_update_kv_lens.call_count, 2) + + +def test_encoder_capture_discards_warmup_and_capture_outputs( + monkeypatch: pytest.MonkeyPatch, +) -> None: discarded = [] @contextlib.contextmanager - def discard(inputs): + def discard(inputs: dict[str, object]) -> Iterator[None]: discarded.append(inputs) yield @@ -51,29 +339,6 @@ def discard(inputs): assert discarded == [capture_inputs] * (runner.WARMUP_STEPS + 1) -def _feature_encoder_runner( - batch_sizes: list[int], fixed_seq_len: int = 1_500 -) -> EncoderCUDAGraphRunner: - """Build a feature-mode backend without enabling CUDA capture.""" - config = EncoderCUDAGraphRunnerConfig( - use_cuda_graph=False, - cuda_graph_padding_enabled=True, - cuda_graph_batch_sizes=batch_sizes, - cuda_graph_num_tokens=[], - cuda_graph_seq_lens=[], - max_cuda_graph_batch_size=max(batch_sizes), - max_cuda_graph_num_tokens=max(batch_sizes) * fixed_seq_len, - max_num_tokens=max(batch_sizes) * fixed_seq_len, - max_seq_len=fixed_seq_len, - is_encoder_decoder=True, - use_fixed_sequence_slots=False, - feature_shape=(480_000,), - feature_dtype=torch.float32, - fixed_seq_len=fixed_seq_len, - ) - return EncoderCUDAGraphRunner(config) - - def test_feature_encoder_capture_keys_are_reachable() -> None: fixed_seq_len = 1_500 batch_sizes = [1, 2, 4, 8] @@ -194,3 +459,30 @@ def test_encoder_decoder_fixed_slots_restore_source_order() -> None: restored_output = runner.restore_encoder_decoder_output(key, fixed_slot_output, inputs) expected_output = torch.cat((fixed_slot_output[511:512], fixed_slot_output[:400])) torch.testing.assert_close(restored_output, expected_output) + + +class SparseGraphKeyTestCase(unittest.TestCase): + def test_sparse_sequence_mode_uses_promoted_context_cursor(self) -> None: + sparse_config = Mock(spec=SeqLenAwareSparseAttentionConfig) + sparse_config.needs_separate_short_long_cuda_graphs.return_value = True + sparse_config.seq_len_threshold = 16 + runner = object.__new__(CUDAGraphRunner) + runner.sparse_config = sparse_config + runner.spec_config = None + runner.graphs = {} + runner.graph_outputs = {} + runner.graph_metadata = {} + runner.padding_dummy_requests = {} + runner.memory_pool = None + + request = _make_request_stub(7, prompt_len=8) + request.py_batch_idx = 0 + request.max_beam_num_tokens = 64 + batch = ScheduledRequests() + batch.generation_requests = [request] + overlap_state = SimpleNamespace(new_tokens=object()) + + self.assertTrue( + runner._get_seq_len_mode(batch, overlap_state, frozenset({request.py_request_id})) + ) + self.assertFalse(runner._get_seq_len_mode(batch, overlap_state, frozenset())) diff --git a/tests/unittest/_torch/executor/test_cuda_graph_runner_gpu.py b/tests/unittest/_torch/executor/test_cuda_graph_runner_gpu.py new file mode 100644 index 000000000000..32838db73702 --- /dev/null +++ b/tests/unittest/_torch/executor/test_cuda_graph_runner_gpu.py @@ -0,0 +1,312 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + + +from __future__ import annotations + +import contextlib +import unittest +from typing import TYPE_CHECKING +from unittest.mock import Mock + +import torch +from _torch.executor.engine_test_utils import DummyModel, _create_request +from _torch.executor.engine_test_utils import kv_cache_manager as managed_kv_cache + +from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata +from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import ( + CUDAGraphRunner, + CUDAGraphRunnerConfig, + KeyType, +) +from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManager, ResourceManagerType +from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests +from tensorrt_llm.mapping import Mapping + +if TYPE_CHECKING: + from collections.abc import Iterator + + from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 + + +@contextlib.contextmanager +def cuda_graph_runner_and_kv_cache( + *, + execution_stream: torch.cuda.Stream | None = None, +) -> Iterator[tuple[CUDAGraphRunner, KVCacheManagerV2]]: + runner = CUDAGraphRunner( + CUDAGraphRunnerConfig( + use_cuda_graph=True, + cuda_graph_padding_enabled=True, + cuda_graph_batch_sizes=[1, 2, 4, 8, 16, 32, 64, 128], + max_cuda_graph_batch_size=128, + max_beam_width=1, + max_num_tokens=258, + spec_config=None, + use_mrope=False, + original_max_total_draft_tokens=0, + enable_attention_dp=False, + is_encoder_decoder=False, + batch_size=13, + mapping=Mapping(), + dist=None, + ) + ) + with managed_kv_cache( + model=DummyModel(torch.half), execution_stream=execution_stream + ) as manager: + try: + yield runner, manager + finally: + runner.clear() + + +class CudaGraphPaddingTestCase(unittest.TestCase): + def test_pad_generation_requests(self) -> None: + (model_engine, kv_cache_manager) = self.enterContext(cuda_graph_runner_and_kv_cache()) + resource_manager = ResourceManager({ResourceManagerType.KV_CACHE_MANAGER: kv_cache_manager}) + + def free_blocks() -> int: + return kv_cache_manager.get_kv_cache_stats().free_num_blocks + + batch_sizes_and_seqlens = [ + (5, 1), + (13, 1), + (5, 25), + ] + for batch_size, max_seq_len in batch_sizes_and_seqlens: + requests = [_create_request(max_seq_len, i) for i in range(batch_size)] + batch = ScheduledRequests() + batch.context_requests_last_chunk = requests + + pages_before = free_blocks() + with model_engine.pad_batch(batch, resource_manager) as padded_batch: + # No padding for prefill + self.assertIs(batch, padded_batch) + self.assertEqual(free_blocks(), pages_before) + + batch = ScheduledRequests() + batch.generation_requests = requests + pages_before = free_blocks() + new_dummy_block = 1 if not model_engine.padding_dummy_requests else 0 + with model_engine.pad_batch(batch, resource_manager) as padded_batch: + if batch_size < 8 and max_seq_len < 25: + self.assertEqual(len(padded_batch.generation_requests) % 8, 0) + else: + # No padding if it would create too many concurrent requests. + # This requirement is not strictly required, but we should probably + # respect the requirement? + # The seqlen check makes sure we don't exceed the KV cache memory + # budget. + self.assertIs(batch, padded_batch) + self.assertEqual(free_blocks() + new_dummy_block, pages_before) + + def test_pad_batch_strips_cudagraph_dummies_on_clean_exit(self) -> None: + # Regression guard for the invariant that CUDAGraphRunner.pad_batch's + # `finally` strips every is_cuda_graph_dummy=True entry from + # scheduled_requests.generation_requests before the `with` block + # exits. Downstream consumers of scheduled_batch.generation_requests + # — including the per-iteration stats populate block in + # PyExecutor._update_iter_stats — rely on never observing + # cudagraph dummies. + (model_engine, kv_cache_manager) = self.enterContext(cuda_graph_runner_and_kv_cache()) + resource_manager = ResourceManager({ResourceManagerType.KV_CACHE_MANAGER: kv_cache_manager}) + + # batch_size=5 rounds up to 8 (nearest captured graph size in the + # fixture config) -> padding_size=3, deterministically. + real_batch_size = 5 + max_seq_len = 1 + real_requests = [_create_request(max_seq_len, i) for i in range(real_batch_size)] + real_ids = [id(r) for r in real_requests] + + batch = ScheduledRequests() + batch.generation_requests = list(real_requests) + + with model_engine.pad_batch(batch, resource_manager) as padded_batch: + # Positive assertion that padding actually fired — guards + # against a vacuous pass where padding was a no-op. + self.assertGreater( + len(padded_batch.generation_requests), + real_batch_size, + "padding did not fire; fixture config may have drifted " + "so that 5 no longer rounds up to 8", + ) + # Every appended entry past the original count is a + # cudagraph-flagged dummy. + for req in padded_batch.generation_requests[real_batch_size:]: + self.assertTrue( + getattr(req, "is_cuda_graph_dummy", False), + "pad_batch appended a request without is_cuda_graph_dummy=True", + ) + # Real requests' identities and order are untouched. + self.assertEqual( + [id(r) for r in padded_batch.generation_requests[:real_batch_size]], real_ids + ) + + # After the with-block: finally must have sliced off the padding. + self.assertEqual( + len(batch.generation_requests), + real_batch_size, + "pad_batch.finally did not strip cudagraph dummies — " + "downstream consumers of scheduled_batch.generation_requests " + "would observe the leaked dummies", + ) + for req in batch.generation_requests: + self.assertFalse( + getattr(req, "is_cuda_graph_dummy", False), + "cudagraph dummy leaked out of pad_batch's finally", + ) + + def test_pad_batch_strips_cudagraph_dummies_on_exception(self) -> None: + # The strip must fire even when the body raises. This is the + # critical property of `finally` vs. a plain trailing statement — + # it guards the invariant on the error path. A refactor that + # accidentally dropped the `finally` would be caught here but not + # by the clean-exit variant. + (model_engine, kv_cache_manager) = self.enterContext(cuda_graph_runner_and_kv_cache()) + resource_manager = ResourceManager({ResourceManagerType.KV_CACHE_MANAGER: kv_cache_manager}) + + real_batch_size = 5 + real_requests = [_create_request(1, i) for i in range(real_batch_size)] + + batch = ScheduledRequests() + batch.generation_requests = list(real_requests) + + class _ForwardBoom(Exception): + pass + + with self.assertRaises(_ForwardBoom): + with model_engine.pad_batch(batch, resource_manager) as padded_batch: + self.assertGreater(len(padded_batch.generation_requests), real_batch_size) + raise _ForwardBoom() + + self.assertEqual(len(batch.generation_requests), real_batch_size) + for req in batch.generation_requests: + self.assertFalse(getattr(req, "is_cuda_graph_dummy", False)) + + def test_preallocate_padding_dummies_uses_captured_draft_lens(self) -> None: + # Speculative engines pad batches with the runtime draft length, so + # the preallocation must create dummies for the draft lengths of the + # captured graphs — not draft_len 0 unconditionally (an extra + # draft_len-0 dummy would permanently hold KV blocks that runtime + # padding never uses). + (model_engine, kv_cache_manager) = self.enterContext(cuda_graph_runner_and_kv_cache()) + resource_manager = ResourceManager({ResourceManagerType.KV_CACHE_MANAGER: kv_cache_manager}) + + runner = model_engine + # Simulate a spec-decode capture: graphs keyed (batch_size, draft_len, + # short_seq_len_mode, is_all_greedy_sample) exist only for draft_len 2. + runner.graphs[(8, 2, False, True)] = Mock() + try: + runner.preallocate_padding_dummies(resource_manager) + + self.assertEqual([2], sorted(runner.padding_dummy_requests.keys())) + finally: + runner.graphs.clear() + for dummy in runner.padding_dummy_requests.values(): + kv_cache_manager.free_resources(dummy) + + def test_release_padding_dummy_covers_every_manager(self) -> None: + # A padding dummy's request ID is spread across several managers, so + # releasing only the main KV cache manager leaves the others holding + # it — and re-creation reuses the same ID. + (model_engine, kv_cache_manager) = self.enterContext(cuda_graph_runner_and_kv_cache()) + spec_manager = Mock() + cross_manager = Mock() + resource_manager = ResourceManager( + { + ResourceManagerType.KV_CACHE_MANAGER: kv_cache_manager, + ResourceManagerType.SPEC_RESOURCE_MANAGER: spec_manager, + ResourceManagerType.CROSS_KV_CACHE_MANAGER: cross_manager, + } + ) + + runner = model_engine + try: + self.assertIsNotNone(runner._get_or_create_padding_dummy(resource_manager, 0)) + dummy = runner.padding_dummy_requests[0] + + self.assertTrue(runner.release_padding_dummy(resource_manager, 0)) + + # Dropped from the runner so the lazy path re-creates it... + self.assertEqual({}, runner.padding_dummy_requests) + # ...and the spec resource manager slot is released too, not just + # the main KV cache manager. + spec_manager.free_resources.assert_called_once_with(dummy) + # The cross-KV manager is only involved for encoder-decoder, which + # this engine is not. + self.assertFalse(runner.is_encoder_decoder) + cross_manager.free_resources.assert_not_called() + + # Releasing again is a no-op rather than a double free. + self.assertFalse(runner.release_padding_dummy(resource_manager, 0)) + spec_manager.free_resources.assert_called_once() + finally: + for dummy in runner.padding_dummy_requests.values(): + kv_cache_manager.free_resources(dummy) + runner.padding_dummy_requests.clear() + + +class ExecutionStreamTestCase(unittest.TestCase): + def test_cuda_graph_replay_observes_execution_stream_dependency(self) -> None: + """KV resume orders graph replay after pending writes on the cache's prior stream.""" + execution_stream = torch.cuda.Stream() + writer_stream = torch.cuda.Stream() + runner, kv_cache_manager = self.enterContext( + cuda_graph_runner_and_kv_cache(execution_stream=execution_stream) + ) + self.addCleanup(execution_stream.synchronize) + self.addCleanup(writer_stream.synchronize) + + request = _create_request(1, 0) + self.assertTrue(kv_cache_manager.prepare_context(request)) + self.assertTrue(kv_cache_manager.resize_context(request, 1)) + cache = kv_cache_manager.kv_cache_map[request.py_request_id] + self.assertGreater(cache.capacity, 0) + + source = kv_cache_manager.get_buffers(0) + assert source is not None + source.zero_() + observed = torch.zeros((), dtype=torch.float32, device="cuda") + key = KeyType(batch_size=1, draft_len=0) + inputs = { + "input_ids": runner.shared_static_tensors["input_ids"][:1], + "position_ids": runner.shared_static_tensors["position_ids"][:, :1], + "attn_metadata": AttentionMetadata( + max_num_requests=1, max_num_tokens=1, kv_cache_manager=kv_cache_manager + ), + } + + def forward(_inputs: dict[str, object]) -> torch.Tensor: + observed.copy_(source.sum(dtype=torch.float32)) + return observed + + torch.cuda.synchronize() + # Warm up lazy CUDA kernels and the event before delaying the real KV writes. + writes_done = torch.cuda.Event() + with torch.cuda.stream(writer_stream): + torch.cuda._sleep(1) + source.fill_(1) + writes_done.record() + writes_done.synchronize() + with torch.cuda.stream(execution_stream): + runner.capture(key, forward, inputs) + torch.cuda.synchronize() + + for value in (7, 11): + cache.cuda_stream = writer_stream.cuda_stream + with torch.cuda.stream(writer_stream): + torch.cuda._sleep(1 << 29) + source.fill_(value) + writes_done.record() + + self.assertFalse(writes_done.query(), "the KV write must still be in flight") + kv_cache_manager.suspend_request(request) + self.assertFalse(cache.is_active) + self.assertTrue(kv_cache_manager.resume_request(request)) + with torch.cuda.stream(execution_stream): + runner.replay(key, inputs) + execution_stream.synchronize() + + self.assertTrue(writes_done.query()) + self.assertEqual(observed.item(), source.numel() * value) diff --git a/tests/unittest/_torch/executor/test_distributed_warmup_oom.py b/tests/unittest/_torch/executor/test_distributed_warmup_oom.py index 6e192b60d528..e9edcacfce4c 100644 --- a/tests/unittest/_torch/executor/test_distributed_warmup_oom.py +++ b/tests/unittest/_torch/executor/test_distributed_warmup_oom.py @@ -1,6 +1,5 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Tests for distributed warmup fail-stop and single-rank recovery.""" import contextlib from collections.abc import Iterator @@ -8,98 +7,12 @@ from unittest import mock import pytest -import torch from tensorrt_llm._torch.pyexecutor import py_executor as py_executor_module -from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import warmup as warmup_module -from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.warmup import ( - DecoderWarmup, - WarmupAgreement, -) -from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManagerType -from tensorrt_llm._torch.pyexecutor.warmup_timer import _WarmupTimer pytestmark = pytest.mark.cpu_only -class _StandInMambaCacheManager: - """Concrete stand-in for the abstract Mamba cache-manager type check.""" - - def __init__(self, available_tokens: int = 8) -> None: - self._available_tokens = available_tokens - - def get_num_available_tokens(self, **_kwargs: object) -> int: - return self._available_tokens - - -def _agreement( - *, - world_size: int = 1, - dwdp_size: int = 0, - tp_size: int = 1, - tp_peer_values: tuple[int, ...] = (), - peer_values: tuple[int, ...] | None = None, - with_dist: bool = True, -) -> WarmupAgreement: - """Build the warmup agreement over a stand-in communicator and mapping. - - ``dist`` is ``Optional`` on the real engine -- stub engines such as - ``DummyModelEngine`` construct without a communicator -- so ``with_dist`` - covers that state too. - - ``tp_allgather`` returns this rank's value followed by ``tp_peer_values``; - ``allgather`` does the same over ``peer_values`` for the forward group. - ``peer_values`` defaults to peers that mirror this rank, i.e. a world that - agrees with whatever this rank reports. - """ - - def _allgather(value): - peers = [value] * (world_size - 1) if peer_values is None else list(peer_values) - return [value, *peers] - - dist = ( - SimpleNamespace( - world_size=world_size, - rank=3, - tp_allgather=lambda value: [value, *tp_peer_values], - allgather=_allgather, - ) - if with_dist - else None - ) - mapping = SimpleNamespace( - dwdp_enabled=dwdp_size > 1, - has_cp_helix=lambda: False, - tp_size=tp_size, - ) - return WarmupAgreement(dist, mapping) - - -def _engine(*, world_size: int = 1, dwdp_size: int = 0) -> DecoderWarmup: - """Build an engine carrying only what the warmup phases read.""" - engine = object.__new__(DecoderWarmup) - engine._agreement = _agreement(world_size=world_size, dwdp_size=dwdp_size) - engine.dist = engine._agreement.dist - engine.mapping = engine._agreement.mapping - engine._reset_moe_alltoall_state = mock.Mock() - engine._config = SimpleNamespace(is_spec_decode=False, max_draft_len=0, spec_config=None) - return engine - - -@contextlib.contextmanager -def _released_batch(batch: object) -> Iterator[object]: - yield batch - - -@contextlib.contextmanager -def _no_cuda_side_effects() -> Iterator[tuple[mock.Mock, mock.Mock]]: - with ( - mock.patch.object(torch.cuda, "empty_cache") as empty_cache, - mock.patch.object(torch.cuda, "synchronize") as synchronize, - ): - yield empty_cache, synchronize - - @contextlib.contextmanager def _guard_env( *, global_size: int, symmetric_crash: bool = False @@ -190,354 +103,3 @@ def test_guard_arms_the_kill_when_the_crash_is_not_proven_symmetric() -> None: probe.assert_called_once_with(4) watchdog.assert_called_once_with(4, error_delivered=None) - - -class TestAllRanksCrashedProbe: - """Local policy of the probe; the Ibarrier itself needs a real world.""" - - def test_single_rank_is_never_symmetric(self) -> None: - from tensorrt_llm._torch.pyexecutor import hang_detector - - assert hang_detector.all_ranks_crashed(1) is False - - def test_mpi_disabled_falls_back_to_the_kill(self) -> None: - from tensorrt_llm._torch.pyexecutor import hang_detector - - with mock.patch.object(hang_detector, "mpi_disabled", return_value=True): - assert hang_detector.all_ranks_crashed(4) is False - - def test_probe_errors_fall_back_to_the_kill(self) -> None: - from mpi4py import MPI - - from tensorrt_llm._torch.pyexecutor import hang_detector - - with ( - mock.patch.object(hang_detector, "mpi_disabled", return_value=False), - mock.patch.object(MPI, "Is_initialized", return_value=True), - mock.patch.object(MPI, "Query_thread", return_value=MPI.THREAD_MULTIPLE), - mock.patch.object(hang_detector, "ENABLE_MULTI_DEVICE", True), - mock.patch.object( - hang_detector, "mpi_comm", side_effect=RuntimeError("no comm") - ) as mpi_comm, - ): - assert hang_detector.all_ranks_crashed(4) is False - mpi_comm.assert_called_once_with() - - def test_disabled_hard_kill_skips_the_barrier(self) -> None: - # With the kill disabled (grace is None) no kill follows a False - # verdict, so posting Ibarrier would leak it into a surviving world. - # The probe must return False without posting the barrier. - from mpi4py import MPI - - from tensorrt_llm._torch.pyexecutor import hang_detector - - comm = mock.Mock() - comm.Get_size.return_value = 4 - with ( - mock.patch.object(hang_detector, "mpi_disabled", return_value=False), - mock.patch.object(MPI, "Is_initialized", return_value=True), - mock.patch.object(MPI, "Query_thread", return_value=MPI.THREAD_MULTIPLE), - mock.patch.object(hang_detector, "ENABLE_MULTI_DEVICE", True), - mock.patch.object(hang_detector, "mpi_comm", return_value=comm), - mock.patch.object(hang_detector, "_rank_crash_kill_grace", return_value=None), - ): - assert hang_detector.all_ranks_crashed(4) is False - comm.Ibarrier.assert_not_called() - - -@pytest.mark.parametrize( - "world_size,dwdp_size,has_batch,expected", - [ - (1, 0, False, False), - # Every rank in the forward group lacks a batch. The allgather proves - # nobody is left inside the collective, so skipping strands no one -- - # this is the shape a max_batch_size=1 config simply cannot build. - (2, 0, False, False), - # DWDP peers live outside `self.dist`, so unanimity cannot be observed - # and the shape stays fatal. - (1, 4, False, None), - (2, 0, True, True), - ], -) -def test_warmup_batch_policy( - world_size: int, dwdp_size: int, has_batch: bool, expected: bool | None -) -> None: - agreement = _agreement(world_size=world_size, dwdp_size=dwdp_size) - batch = object() if has_batch else None - if expected is None: - with pytest.raises(RuntimeError, match="cannot skip the shape"): - agreement.should_run_warmup_batch(batch, 128, "general") - else: - assert agreement.should_run_warmup_batch(batch, 128, "general") is expected - - -def test_symmetric_absence_skips_instead_of_failing_the_world() -> None: - """The regression this fix exists for. - - `max_batch_size=1` cannot build the mixed context+generation attention - shape, on any rank, at any KV cache size. Every rank reaches the same - verdict from configuration alone, so the group agrees and skips. Failing - here took down ten multi-GPU tests that had always skipped this shape. - """ - agreement = _agreement(world_size=2, tp_size=2, peer_values=(0,)) - - assert ( - agreement.should_run_warmup_batch(None, 2, "attention, num_tokens=2, num_gen_requests=1") - is False - ) - - -def test_group_disagreement_still_fails_with_the_ranks_named() -> None: - """Asymmetric capacity is the deadlock the agreement exists to catch.""" - agreement = _agreement(world_size=3, peer_values=(1, 0)) - - with pytest.raises(RuntimeError) as excinfo: - agreement.should_run_warmup_batch(None, 128, "general") - - message = str(excinfo.value) - assert "rank(s) [0, 2]" in message - assert "would deadlock" in message - - -def test_pipeline_only_world_agrees_without_a_tp_group() -> None: - """TP=1 with pipeline peers is where the TP-only check gathered nothing. - - Three of the ten failures were disaggregated context servers at - `tensor_parallel_size: 1, pipeline_parallel_size: 2`. - """ - agreement = _agreement(world_size=2, tp_size=1, peer_values=(0,)) - - assert agreement.should_run_warmup_batch(None, 128, "general") is False - - -def test_agreement_group_is_absent_for_dwdp() -> None: - agreement = _agreement(world_size=4, dwdp_size=4) - - assert agreement.allgather() is None - - -@pytest.mark.parametrize( - "world_size,peer_values,flag,expected", - [ - (1, None, True, True), - (2, (1,), True, True), - # One pipeline stage owns the guided decoder and opts out; the group - # follows it rather than letting the others enter alone. - (2, (0,), True, False), - (2, (1,), False, False), - ], -) -def test_phase_entry_flag_is_agreed( - world_size: int, peer_values: tuple[int, ...] | None, flag: bool, expected: bool -) -> None: - agreement = _agreement(world_size=world_size, peer_values=peer_values) - - assert agreement.agree_flag(flag) is expected - - -def test_shape_list_narrows_to_what_every_rank_proposed() -> None: - """Attention-DP gives ranks different capacities and so different shapes.""" - agreement = _agreement(world_size=2) - agreement.dist.allgather = lambda value: [value, [[1, 0], [2, 0]]] - - assert agreement.agree_shapes([(1, 0), (2, 0), (4096, 0)]) == [(1, 0), (2, 0)] - - -def test_shape_list_is_untouched_without_a_group() -> None: - agreement = _agreement(world_size=1) - configs = [(1, 0), (128, 0)] - - assert agreement.agree_shapes(configs) == configs - - -def test_warmup_batch_policy_without_communicator() -> None: - """An engine built without a communicator cannot strand anyone: skip, don't raise.""" - agreement = _agreement(with_dist=False) - - assert agreement.should_run_warmup_batch(None, 128, "general") is False - - -def test_tp_agreement_needs_a_communicator_not_just_a_mapping() -> None: - """A mapping can claim TP peers that no communicator exists to reach. - - The agreement check gates on `tp_size`, so it has to answer for that - state rather than walking into a collective on `None`. - """ - agreement = _agreement(tp_size=2, with_dist=False) - - assert agreement.should_run_warmup_batch(object(), 128, "general") is True - assert agreement.should_run_warmup_batch(None, 128, "general") is False - - -def test_tp_disagreement_reports_the_ranks_that_lost_their_batch() -> None: - """Asymmetric attention-DP capacity is the deadlock this check exists for. - - The rank that still has a batch must not walk into a forward its peer - will never join, and the failure has to name the peer so the KV cache - fraction can be raised. - - DWDP is what still reaches the TP-only check: its peers are outside - ``self.dist``, so no group agreement is available and the narrower TP - check remains the only thing standing between an asymmetric world and a - deadlock. - """ - agreement = _agreement(world_size=2, dwdp_size=4, tp_size=2, tp_peer_values=(0,)) - - with pytest.raises(RuntimeError) as excinfo: - agreement.should_run_warmup_batch(object(), 128, "general") - - message = str(excinfo.value) - assert "TP rank(s) [1]" in message - assert "[128, 0]" in message - - -def test_tp_agreement_lets_a_symmetric_world_run() -> None: - agreement = _agreement(world_size=2, tp_size=2, tp_peer_values=(1,)) - - assert agreement.should_run_warmup_batch(object(), 128, "general") is True - - -def _general_warmup_engine(*, world_size: int, dwdp_size: int) -> DecoderWarmup: - engine = _engine(world_size=world_size, dwdp_size=dwdp_size) - engine._warmup_timer = _WarmupTimer(rank=engine.dist.rank) - batch = object() - engine._dummy_requests = SimpleNamespace( - create_warmup_request=mock.Mock(return_value=batch), - release_batch_context=lambda *_a, **_kw: _released_batch(batch), - ) - return engine - - -@pytest.mark.parametrize( - "world_size,dwdp_size,is_fatal", - [(1, 0, False), (2, 0, True), (1, 4, True)], -) -def test_general_warmup_oom_policy(world_size: int, dwdp_size: int, is_fatal: bool) -> None: - engine = _general_warmup_engine(world_size=world_size, dwdp_size=dwdp_size) - error = torch.OutOfMemoryError("asymmetric OOM") - engine._forward_warmup = mock.Mock(side_effect=error if is_fatal else [error, None]) - - with _no_cuda_side_effects() as (empty_cache, _synchronize): - if is_fatal: - with pytest.raises(torch.OutOfMemoryError) as excinfo: - engine._general_warmup_impl(object(), [(128, 0), (64, 0)]) - assert excinfo.value is error - else: - engine._general_warmup_impl(object(), [(128, 0), (64, 0)]) - - if is_fatal: - # The remaining shape is never attempted: this rank's peers are - # already stuck in the failed forward's collectives. - engine._forward_warmup.assert_called_once() - else: - assert engine._forward_warmup.call_count == 2 - # A retry after an OOM between dispatch() and combine() has to start - # from a clean MoE all-to-all state. - engine._reset_moe_alltoall_state.assert_called_once_with() - empty_cache.assert_called_once_with() - - -_KV_ALLOC_ERROR = "Can't allocate new blocks for window size 8" - - -def _mamba_engine(*, world_size: int = 1, dwdp_size: int = 0) -> tuple[DecoderWarmup, object]: - engine = _engine(world_size=world_size, dwdp_size=dwdp_size) - engine._config = SimpleNamespace( - **vars(engine._config), - max_num_tokens=8, - max_batch_size=4, - max_seq_len=8, - original_max_draft_len=0, - enable_autotuner=False, - ) - engine._graphs = SimpleNamespace(no_cuda_graph=contextlib.nullcontext) - # The warmup resolves its chunk-alignment variant off the model's Mamba - # metadata class; a model declaring none takes the ``Mamba2Metadata`` - # default, i.e. no alignment rewrite. - engine.model = SimpleNamespace() - batch = object() - engine._dummy_requests = SimpleNamespace( - release_batch_context=lambda *_a, **_kw: _released_batch(batch) - ) - - cache_manager = _StandInMambaCacheManager() - resource_manager = SimpleNamespace( - get_resource_manager=lambda key: ( - cache_manager if key == ResourceManagerType.KV_CACHE_MANAGER else None - ) - ) - return engine, resource_manager - - -def _run_mamba_warmup(engine: DecoderWarmup, resource_manager: object) -> None: - with ( - mock.patch.object(warmup_module, "MambaHybridCacheManager", _StandInMambaCacheManager), - mock.patch.object( - warmup_module.Mamba2Metadata, - "force_initial_states_for_warmup", - side_effect=contextlib.nullcontext, - ), - mock.patch.object(warmup_module, "clear_memory_buffers"), - ): - engine._run_mamba_hybrid_warmup(resource_manager) - - -@pytest.mark.parametrize( - "error,recoverable", - [ - (torch.OutOfMemoryError("OOM"), True), - (RuntimeError(_KV_ALLOC_ERROR), True), - (RuntimeError("unexpected"), False), - ], -) -def test_mamba_preforward_error_policy_when_alone(error: Exception, recoverable: bool) -> None: - engine, resource_manager = _mamba_engine() - engine._dummy_requests.create_warmup_request = mock.Mock(side_effect=error) - engine._forward_warmup = mock.Mock() - - with _no_cuda_side_effects(): - if recoverable: - _run_mamba_warmup(engine, resource_manager) - else: - with pytest.raises(RuntimeError, match="unexpected"): - _run_mamba_warmup(engine, resource_manager) - - engine._forward_warmup.assert_not_called() - # The failure predates dispatch(), so there is no half-finished MoE - # all-to-all exchange to unwind. - engine._reset_moe_alltoall_state.assert_not_called() - - -def test_mamba_midforward_runtime_error_recovers_when_alone() -> None: - engine, resource_manager = _mamba_engine() - engine._dummy_requests.create_warmup_request = mock.Mock(return_value=object()) - engine._forward_warmup = mock.Mock(side_effect=RuntimeError("mid-forward failure")) - - with _no_cuda_side_effects(): - _run_mamba_warmup(engine, resource_manager) - - # Every shape is attempted, and each failed forward has to leave the MoE - # all-to-all state clean for the shape that follows it. - assert engine._forward_warmup.call_count >= 1 - assert engine._reset_moe_alltoall_state.call_count == engine._forward_warmup.call_count - - -@pytest.mark.parametrize("phase", ["pre-forward", "mid-forward"]) -def test_mamba_error_is_fatal_when_distributed(phase: str) -> None: - engine, resource_manager = _mamba_engine(world_size=2) - error = RuntimeError(_KV_ALLOC_ERROR) - if phase == "pre-forward": - engine._dummy_requests.create_warmup_request = mock.Mock(side_effect=error) - engine._forward_warmup = mock.Mock() - else: - engine._dummy_requests.create_warmup_request = mock.Mock(return_value=object()) - engine._forward_warmup = mock.Mock(side_effect=error) - - with _no_cuda_side_effects(): - with pytest.raises(RuntimeError) as excinfo: - _run_mamba_warmup(engine, resource_manager) - - assert excinfo.value is error - # The remaining shape is never attempted: recovering locally would leave - # peers waiting in a forward this rank has abandoned. - engine._dummy_requests.create_warmup_request.assert_called_once() diff --git a/tests/unittest/_torch/executor/test_hang_detector.py b/tests/unittest/_torch/executor/test_hang_detector.py new file mode 100644 index 000000000000..a5322340464f --- /dev/null +++ b/tests/unittest/_torch/executor/test_hang_detector.py @@ -0,0 +1,58 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from unittest import mock + +import pytest + + +@pytest.mark.cpu_only +class TestAllRanksCrashedProbe: + def test_single_rank_is_never_symmetric(self) -> None: + from tensorrt_llm._torch.pyexecutor import hang_detector + + assert hang_detector.all_ranks_crashed(1) is False + + def test_mpi_disabled_falls_back_to_the_kill(self) -> None: + from tensorrt_llm._torch.pyexecutor import hang_detector + + with mock.patch.object(hang_detector, "mpi_disabled", return_value=True): + assert hang_detector.all_ranks_crashed(4) is False + + def test_probe_errors_fall_back_to_the_kill(self) -> None: + from mpi4py import MPI + + from tensorrt_llm._torch.pyexecutor import hang_detector + + with ( + mock.patch.object(hang_detector, "mpi_disabled", return_value=False), + mock.patch.object(MPI, "Is_initialized", return_value=True), + mock.patch.object(MPI, "Query_thread", return_value=MPI.THREAD_MULTIPLE), + mock.patch.object(hang_detector, "ENABLE_MULTI_DEVICE", True), + mock.patch.object( + hang_detector, "mpi_comm", side_effect=RuntimeError("no comm") + ) as mpi_comm, + ): + assert hang_detector.all_ranks_crashed(4) is False + mpi_comm.assert_called_once_with() + + def test_disabled_hard_kill_skips_the_barrier(self) -> None: + # With the kill disabled (grace is None) no kill follows a False + # verdict, so posting Ibarrier would leak it into a surviving world. + # The probe must return False without posting the barrier. + from mpi4py import MPI + + from tensorrt_llm._torch.pyexecutor import hang_detector + + comm = mock.Mock() + comm.Get_size.return_value = 4 + with ( + mock.patch.object(hang_detector, "mpi_disabled", return_value=False), + mock.patch.object(MPI, "Is_initialized", return_value=True), + mock.patch.object(MPI, "Query_thread", return_value=MPI.THREAD_MULTIPLE), + mock.patch.object(hang_detector, "ENABLE_MULTI_DEVICE", True), + mock.patch.object(hang_detector, "mpi_comm", return_value=comm), + mock.patch.object(hang_detector, "_rank_crash_kill_grace", return_value=None), + ): + assert hang_detector.all_ranks_crashed(4) is False + comm.Ibarrier.assert_not_called() diff --git a/tests/unittest/_torch/executor/test_model_engine.py b/tests/unittest/_torch/executor/test_model_engine.py new file mode 100644 index 000000000000..c472dbf40d2f --- /dev/null +++ b/tests/unittest/_torch/executor/test_model_engine.py @@ -0,0 +1,1443 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + + +from __future__ import annotations + +import gc +import unittest +import weakref +from ast import literal_eval +from contextlib import nullcontext +from types import SimpleNamespace +from typing import TYPE_CHECKING +from unittest.mock import Mock, call, patch + +import pytest +import torch +from _torch.executor.engine_test_utils import ( + DummyModel, + DummyModelEngine, + _AttentionBackend, + _clone_tensors, + _create_request, + _create_request_with_tokens, + _model_engine_with_runner, + create_model_engine_and_kvcache, +) + +import tensorrt_llm._torch.pyexecutor.model_engine as model_engine_module +from tensorrt_llm._torch.attention.backends.interface import AttentionRuntimeFeatures +from tensorrt_llm._torch.models.modeling_multimodal_mixin import MultimodalModelMixin +from tensorrt_llm._torch.models.modeling_utils import DecoderModelForCausalLM +from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import KeyType +from tensorrt_llm._torch.pyexecutor.engine.model_call import ModelCaller +from tensorrt_llm._torch.pyexecutor.engine.multimodal import is_multimodal +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import DecoderRunner +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.graphs import DecoderCudaGraphs +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.inputs import ( + DecoderInputBuffers, + DecoderInputPreparer, +) +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.runner import PrefillCompiledModel +from tensorrt_llm._torch.pyexecutor.engine.runners.encoder_decoder import EncoderDecoderRunner +from tensorrt_llm._torch.pyexecutor.engine.runners.no_kv_cache import NoKVCacheRunner +from tensorrt_llm._torch.pyexecutor.engine.runners.pooling import PoolingRunner +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest +from tensorrt_llm._torch.pyexecutor.model_engine import ModelEngine +from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManager, ResourceManagerType +from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests +from tensorrt_llm._torch.utils import torch_compiling +from tensorrt_llm.inputs.registry import BaseMultimodalInputProcessor +from tensorrt_llm.llmapi import CudaGraphConfig, KvCacheConfig, SADecodingConfig +from tensorrt_llm.llmapi.llm_args import ( + EncodeCudaGraphConfig, + MultimodalConfig, + PrefillCudaGraphBackend, + TorchLlmArgs, +) +from tensorrt_llm.mapping import Mapping + +if TYPE_CHECKING: + from typing import Any + + from tensorrt_llm._torch.pyexecutor.sampler import SampleStateTensors + + +class DummyLegacyMultimodalIndexModel(MultimodalModelMixin, torch.nn.Module): + class Config: + vocab_size = 100 + + config = Config() + + @property + def mm_token_ids(self) -> torch.Tensor: + return torch.tensor([90, 91], dtype=torch.int32) + + +def _prepare_kv_cache_resources(manager: KVCacheManagerV2, batch: ScheduledRequests) -> None: + # These model-engine tests bypass the scheduler, which owns V2 allocation. + for request in batch.context_requests: + assert manager.prepare_context(request) + assert manager.resize_context( + request, request.prompt_len - request.context_current_position + ) + for request in batch.generation_requests: + assert manager.try_allocate_generation(request) + manager.prepare_resources(batch) + + +def _install_graph_runner_stub(runner: DecoderRunner, key: KeyType | None = None) -> Mock: + """Replace the runner's CUDA graph runner with one that never pads. + + Every lookup hits `stub.key`, or misses while it is None. A hit replays by + recording a copy of the prepared inputs. + """ + stub = Mock() + stub.enabled = True + stub.is_warmup_only = False + stub.key = key + stub.pad_batch.side_effect = lambda batch, *_args: nullcontext(batch) + stub.get_graph_pool.return_value = None + stub.needs_capture.return_value = False + stub.maybe_get_cuda_graph.side_effect = lambda batch, **kwargs: ( + (kwargs["attn_metadata"], kwargs["spec_metadata"], stub.key) + if stub.key is not None + else (None, None, None) + ) + stub.replayed_inputs = [] + + def replay( + key: KeyType, inputs: dict[str, Any], position_id_offsets: torch.Tensor | None = None + ) -> dict[str, torch.Tensor]: + stub.replayed_inputs.append(_clone_tensors(inputs)) + return {"logits": torch.zeros((inputs["input_ids"].shape[0], 10), device="cuda")} + + stub.replay.side_effect = replay + runner._graphs.cuda_graph_runner = stub + return stub + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("failure_stage", [None, "loader", "graphs", "gc"]) +def test_cleanup_releases_runner_model_before_release_gc( + monkeypatch: pytest.MonkeyPatch, + failure_stage: str | None, +) -> None: + engine = object.__new__(ModelEngine) + model = torch.nn.Linear(1, 1) + model_ref = weakref.ref(model) + error = RuntimeError("release failed") + loader = Mock() + if failure_stage == "loader": + loader.cleanup.side_effect = [error, None] + release_graphs = Mock(side_effect=[error, None] if failure_stage == "graphs" else None) + engine._cleanup_done = False + engine.model_loader = loader + engine._runner = SimpleNamespace(_model=model, release_graphs=release_graphs) + engine._torch_compile_backend = None + engine._model_caller = ModelCaller(model) + engine._mm_item_scheduler = None + engine._lora_model = model + engine._lora_owner = None + engine.model = model + engine.input_processor = None + del model + + def collect() -> None: + assert model_ref() is None + if failure_stage == "gc" and release_gc.call_count == 1: + raise error + + release_gc = Mock(side_effect=collect) + monkeypatch.setattr(model_engine_module, "release_gc", release_gc) + if failure_stage is not None: + with pytest.raises(RuntimeError, match="release failed"): + engine.cleanup() + assert not engine._cleanup_done + if failure_stage != "gc": + assert model_ref() is not None + assert engine.model_loader is (loader if failure_stage == "loader" else None) + engine.cleanup() + engine.cleanup() + assert engine._cleanup_done + assert model_ref() is None + assert loader.cleanup.call_count == (2 if failure_stage == "loader" else 1) + assert release_graphs.call_count == (2 if failure_stage == "graphs" else 1) + release_graphs.assert_called_with(release_nccl_window_owners=False) + assert release_gc.call_count == (2 if failure_stage == "gc" else 1) + + +class GraphOwnerTeardownTestCase(unittest.TestCase): + def test_engine_graph_ownership_is_split(self) -> None: + clear_blocks_graph = Mock() + model = Mock() + model.modules.return_value = [SimpleNamespace(clear_blocks_cuda_graph=clear_blocks_graph)] + backend = Mock() + cuda_runner = Mock() + breakable_runner = Mock() + decoder_runner = object.__new__(DecoderRunner) + decoder_runner._graphs = object.__new__(DecoderCudaGraphs) + decoder_runner._graphs.cuda_graph_runner = cuda_runner + decoder_runner._graphs.breakable_cuda_graph_runner = breakable_runner + engine = SimpleNamespace( + model=model, + _torch_compile_backend=backend, + _runner=decoder_runner, + ) + + ModelEngine._release_cuda_graphs(engine, release_nccl_window_owners=False) + + clear_blocks_graph.assert_not_called() + backend.clear_piecewise_cuda_graphs.assert_called_once_with( + release_nccl_window_owners=False + ) + for runner in (cuda_runner, breakable_runner): + runner.clear.assert_called_once_with(release_nccl_window_owners=False) + + ModelEngine._release_model_owned_cuda_graphs(engine, release_nccl_window_owners=False) + + clear_blocks_graph.assert_called_once_with(release_nccl_window_owners=False) + + def test_executor_release_clears_model_owned_graphs_first(self) -> None: + for release_model_owned_graphs in (True, False): + with self.subTest(release_model_owned_graphs=release_model_owned_graphs): + engine = Mock() + + ModelEngine.release_executor_resources( + engine, release_model_owned_graphs=release_model_owned_graphs + ) + + expected = [call._release_cuda_graphs()] + if release_model_owned_graphs: + expected.insert(0, call._release_model_owned_cuda_graphs()) + self.assertEqual(engine.method_calls, expected) + + +class DecoderForwardTestCase(unittest.TestCase): + def _make_engine( + self, + model: DummyModel | None = None, + ) -> tuple[ModelEngine, KVCacheManagerV2, ResourceManager, Mock]: + (engine, kv_cache_manager) = self.enterContext(create_model_engine_and_kvcache(model=model)) + resource_manager = ResourceManager({ResourceManagerType.KV_CACHE_MANAGER: kv_cache_manager}) + graph_runner = _install_graph_runner_stub(engine._runner) + return engine, kv_cache_manager, resource_manager, graph_runner + + @staticmethod + def _forward( + engine: ModelEngine, + kv_cache_manager: KVCacheManagerV2, + resource_manager: ResourceManager, + batch: ScheduledRequests, + new_tensors_device: SampleStateTensors | None = None, + ) -> dict[str, torch.Tensor | int]: + _prepare_kv_cache_resources(kv_cache_manager, batch) + return engine.forward(batch, resource_manager, new_tensors_device) + + @staticmethod + def _request(tokens: list[int], req_id: int, seq_slot: int) -> LlmRequest: + request = _create_request_with_tokens(tokens, req_id) + request.py_seq_slot = seq_slot + return request + + def _prefill( + self, + engine: ModelEngine, + kv_cache_manager: KVCacheManagerV2, + resource_manager: ResourceManager, + *requests: LlmRequest, + ) -> None: + batch = ScheduledRequests() + batch.context_requests_last_chunk = list(requests) + self._forward(engine, kv_cache_manager, resource_manager, batch) + + def _prefill_all_but_last_token( + self, + engine: ModelEngine, + kv_cache_manager: KVCacheManagerV2, + resource_manager: ResourceManager, + request: LlmRequest, + ) -> None: + """Leave `request` with a final one-token context chunk.""" + batch = ScheduledRequests() + batch.context_requests_chunking = [request] + _prepare_kv_cache_resources(kv_cache_manager, batch) + # Like the scheduler, narrow the chunk after V2 spans it to the prompt end. + request.context_chunk_size = request.py_prompt_len - 1 + engine.forward(batch, resource_manager) + request.move_to_next_context_chunk() + request.context_chunk_size = 1 + + def _final_context_and_generation( + self, + engine: ModelEngine, + kv_cache_manager: KVCacheManagerV2, + resource_manager: ResourceManager, + ) -> tuple[LlmRequest, LlmRequest, ScheduledRequests]: + context = self._request([11, 22, 33, 44], 1, seq_slot=0) + generation = self._request([50, 51, 52, 53, 54], 2, seq_slot=1) + self._prefill_all_but_last_token(engine, kv_cache_manager, resource_manager, context) + self._prefill(engine, kv_cache_manager, resource_manager, generation) + generation.add_new_token(60, 0) + batch = ScheduledRequests() + batch.context_requests_last_chunk = [context] + batch.generation_requests = [generation] + return context, generation, batch + + def test_forward_commits_candidate_only_on_graph_hit(self) -> None: + engine, kv_cache_manager, resource_manager, graph_runner = self._make_engine() + context, generation, batch = self._final_context_and_generation( + engine, kv_cache_manager, resource_manager + ) + num_model_calls = len(engine.model.recorded_kwargs) + graph_runner.key = KeyType(batch_size=2, draft_len=0) + + outputs = self._forward(engine, kv_cache_manager, resource_manager, batch) + + lookup = graph_runner.maybe_get_cuda_graph.call_args + self.assertIsNot(lookup.args[0], batch) + self.assertEqual(lookup.args[0].generation_requests, [context, generation]) + self.assertEqual( + lookup.kwargs["promoted_context_request_ids"], frozenset({context.py_request_id}) + ) + graph_runner.replay.assert_called_once() + self.assertEqual(graph_runner.replay.call_args.args[0], graph_runner.key) + replayed = graph_runner.replayed_inputs[0] + self.assertEqual(replayed["input_ids"].tolist(), [44, 60]) + self.assertEqual(replayed["position_ids"].flatten().tolist(), [3, 5]) + self.assertEqual(len(engine.model.recorded_kwargs), num_model_calls) + self.assertEqual(set(outputs), {"logits", "runtime_draft_len"}) + self.assertEqual(engine.attn_metadata.num_contexts, 0) + self.assertEqual(engine.iter_states["num_ctx_requests"], 1) + self.assertEqual(engine.iter_states["num_ctx_tokens"], 1) + self.assertEqual(engine.iter_states["num_generation_tokens"], 1) + + def test_forward_graph_miss_uses_semantic_eager_batch(self) -> None: + engine, kv_cache_manager, resource_manager, graph_runner = self._make_engine() + context, generation, batch = self._final_context_and_generation( + engine, kv_cache_manager, resource_manager + ) + + self._forward(engine, kv_cache_manager, resource_manager, batch) + + lookup = graph_runner.maybe_get_cuda_graph.call_args + self.assertEqual(lookup.args[0].generation_requests, [context, generation]) + graph_runner.replay.assert_not_called() + model_kwargs = engine.model.recorded_kwargs[-1] + self.assertEqual(model_kwargs["input_ids"].tolist(), [44, 60]) + self.assertEqual(model_kwargs["position_ids"].flatten().tolist(), [3, 5]) + self.assertEqual(engine.attn_metadata.num_contexts, 1) + + def test_forward_allows_guided_context_logits_on_graph_hit(self) -> None: + engine, kv_cache_manager, resource_manager, graph_runner = self._make_engine() + engine._runner.guided_decoder = Mock() + context = self._request([11, 22, 33, 44], 1, seq_slot=0) + self._prefill_all_but_last_token(engine, kv_cache_manager, resource_manager, context) + context.py_return_context_logits = True + batch = ScheduledRequests() + batch.context_requests_last_chunk = [context] + graph_runner.key = KeyType(batch_size=1, draft_len=0) + + self._forward(engine, kv_cache_manager, resource_manager, batch) + + lookup = graph_runner.maybe_get_cuda_graph.call_args + self.assertEqual(lookup.args[0].generation_requests, [context]) + self.assertEqual( + lookup.kwargs["promoted_context_request_ids"], frozenset({context.py_request_id}) + ) + graph_runner.replay.assert_called_once() + self.assertEqual(graph_runner.replayed_inputs[0]["input_ids"].tolist(), [44]) + + def test_promoted_context_uses_prompt_token_during_overlap(self) -> None: + engine, kv_cache_manager, resource_manager, graph_runner = self._make_engine() + context = self._request([11, 22, 33, 44], 1, seq_slot=0) + generation = self._request([50, 51, 52, 53, 54], 2, seq_slot=1) + self._prefill_all_but_last_token(engine, kv_cache_manager, resource_manager, context) + self._prefill(engine, kv_cache_manager, resource_manager, generation) + # The chunked prefill stands in for a reused prefix, so the promoted + # row is the context's first forward and latches its cached tokens. + context._cached_tokens_set = False + # A promoted context row must ignore any stale overlap slot. + context.py_batch_idx = 3 + context.py_num_compressed_tokens = 1 + generation.py_batch_idx = 1 + new_tokens = torch.zeros((1, 4, 1), dtype=torch.int32, device="cuda") + new_tokens[0, 0, 0] = 999 + new_tokens[0, 1, 0] = 777 + batch = ScheduledRequests() + batch.context_requests_last_chunk = [context] + batch.generation_requests = [generation] + graph_runner.key = KeyType(batch_size=2, draft_len=0) + + self._forward( + engine, + kv_cache_manager, + resource_manager, + batch, + SimpleNamespace(new_tokens=new_tokens), + ) + + replayed = graph_runner.replayed_inputs[0] + self.assertEqual(replayed["input_ids"].tolist(), [44, 777]) + self.assertEqual(replayed["position_ids"].flatten().tolist(), [3, 5]) + attn_metadata = engine.attn_metadata + self.assertEqual(attn_metadata.num_contexts, 0) + self.assertEqual(attn_metadata.kv_cache_params.num_cached_tokens_per_seq, [2, 5]) + self.assertEqual(context.cached_tokens, 3) + + def test_steady_overlap_generation_advances_positions(self) -> None: + engine, kv_cache_manager, resource_manager, _ = self._make_engine() + requests = [ + self._request([10, 11, 12, 13], 1, seq_slot=0), + self._request([20, 21, 22, 23, 24, 25], 2, seq_slot=1), + ] + self._prefill(engine, kv_cache_manager, resource_manager, *requests) + for batch_idx, request in enumerate(requests): + request.py_batch_idx = batch_idx + batch = ScheduledRequests() + batch.generation_requests = requests + steps = ((100, 200), (101, 201), (102, 202)) + + for step, tokens in enumerate(steps): + new_tokens = torch.tensor(tokens, dtype=torch.int32, device="cuda").view(1, 2, 1) + self._forward( + engine, + kv_cache_manager, + resource_manager, + batch, + SimpleNamespace(new_tokens=new_tokens), + ) + model_kwargs = engine.model.recorded_kwargs[-1] + with self.subTest(step=step): + self.assertEqual(model_kwargs["input_ids"].tolist(), list(tokens)) + self.assertEqual( + model_kwargs["position_ids"].flatten().tolist(), [4 + step, 6 + step] + ) + # The sampler appends the previous step's token after the next + # step has been scheduled. + for request, token in zip(requests, tokens): + request.add_new_token(token, 0) + + def test_forward_callback_observes_eager_and_replay_execution(self) -> None: + for graph_hit in (False, True): + with self.subTest(graph_hit=graph_hit): + engine, manager, resources, graphs = self._make_engine() + request = self._request([11, 22, 33], 1, 0) + self._prefill(engine, manager, resources, request) + request.add_new_token(44, 0) + batch = ScheduledRequests() + batch.generation_requests = [request] + graphs.key = KeyType(batch_size=1, draft_len=0) if graph_hit else None + observed = [] + + def after_forward() -> None: + inputs = ( + graphs.replayed_inputs[-1] + if graph_hit + else engine.model.recorded_kwargs[-1] + ) + observed.append(inputs["input_ids"].tolist()) + + engine.register_forward_pass_callable(after_forward) + self._forward(engine, manager, resources, batch) + self.assertEqual(observed, [[44]]) + self.assertEqual(graphs.replay.call_count, int(graph_hit)) + + def _make_mrope_engine(self) -> tuple[ModelEngine, KVCacheManagerV2, ResourceManager, Mock]: + return self._make_engine(DummyModel(torch.half, rope_scaling={"type": "mrope"})) + + def _mrope_generation( + self, + engine: ModelEngine, + kv_cache_manager: KVCacheManagerV2, + resource_manager: ResourceManager, + num_tokens: int, + req_id: int, + seq_slot: int, + delta: int | None, + ) -> LlmRequest: + """A generation request at position `num_tokens - 1`. + + It carries an MRoPE delta when `delta` is not None, and no + multimodal data at all otherwise (a text-only prompt). + """ + request = self._request(list(range(1, num_tokens)), req_id, seq_slot) + self._prefill(engine, kv_cache_manager, resource_manager, request) + request.add_new_token(num_tokens, 0) + if delta is None: + request.py_multimodal_data = {} + else: + request.py_multimodal_data = { + "mrope_config": { + "mrope_position_deltas": torch.tensor([[delta]], dtype=torch.int32) + }, + } + return request + + def _forward_generation( + self, + engine: ModelEngine, + kv_cache_manager: KVCacheManagerV2, + resource_manager: ResourceManager, + requests: list[LlmRequest], + ) -> dict[str, Any]: + batch = ScheduledRequests() + batch.generation_requests = requests + self._forward(engine, kv_cache_manager, resource_manager, batch) + return engine.model.recorded_kwargs[-1] + + def test_mrope_generation_with_partial_segments(self) -> None: + """A real multimodal span next to a padding dummy without MRoPE data.""" + engine, kv_cache_manager, resource_manager, _ = self._make_mrope_engine() + multimodal = self._mrope_generation(engine, kv_cache_manager, resource_manager, 4, 1, 0, 10) + multimodal.py_multimodal_data["multimodal_embedding"] = torch.ones( + (1, 1), dtype=torch.float16 + ) + # A generation-only admission has no batch slot yet, which is when + # the context-phase payload is stripped. + multimodal.py_batch_idx = None + dummy = self._mrope_generation(engine, kv_cache_manager, resource_manager, 6, 2, 1, None) + dummy.is_cuda_graph_dummy = True + + model_kwargs = self._forward_generation( + engine, kv_cache_manager, resource_manager, [multimodal, dummy] + ) + + expected = torch.tensor([[[13, 5]]] * 3, dtype=torch.int32, device="cuda") + torch.testing.assert_close(model_kwargs["position_ids"], expected, atol=0, rtol=0) + self.assertEqual(model_kwargs["mrope_delta_write_seq_slots"].cpu().tolist(), [0]) + # Read slots are dense w.r.t. the generation batch: the padded dummy + # has no MRoPE metadata, so it resolves to the reserved zero slot + # (max_num_tokens * pp_size) rather than being dropped, which would + # shift every later request onto another request's delta. + self.assertEqual( + model_kwargs["mrope_delta_read_seq_slots"].cpu().tolist(), [0, engine.max_num_tokens] + ) + self.assertNotIn("multimodal_embedding", multimodal.py_multimodal_data) + + def test_mrope_mixed_text_only_keeps_deltas_dense(self) -> None: + """A text-only request between two multimodal ones must not compact the + MRoPE delta read slots. + + The attention kernel indexes `mrope_position_deltas` by generation batch + index, so a list that skips the text-only request would hand request 2's + delta to the text-only request and read out of bounds for request 2. + """ + engine, kv_cache_manager, resource_manager, _ = self._make_mrope_engine() + # (num_tokens, req_id, seq_slot, delta); the middle one is text-only. + requests = [ + self._mrope_generation(engine, kv_cache_manager, resource_manager, *args) + for args in ((4, 1, 0, 10), (5, 2, 1, None), (6, 3, 2, 20)) + ] + + model_kwargs = self._forward_generation( + engine, kv_cache_manager, resource_manager, requests + ) + + # One entry per generation request, in batch order; the reserved zero + # slot stands in for the text-only request's zero delta. + self.assertEqual( + model_kwargs["mrope_delta_read_seq_slots"].cpu().tolist(), [0, engine.max_num_tokens, 2] + ) + # Only the two multimodal requests seed the seq-slot delta cache. + self.assertEqual(model_kwargs["mrope_delta_write_seq_slots"].cpu().tolist(), [0, 2]) + # The position is num_tokens - 1, offset by the request's delta; the + # text-only request keeps the plain scalar position on all 3 axes. + expected = torch.tensor([[[13, 4, 25]]] * 3, dtype=torch.int32, device="cuda") + torch.testing.assert_close(model_kwargs["position_ids"], expected, atol=0, rtol=0) + + def test_mrope_all_text_only_drops_deltas(self) -> None: + """A generation batch with no MRoPE metadata at all emits no delta + tensors, so the steady-state generation fast path stays reachable.""" + engine, kv_cache_manager, resource_manager, _ = self._make_mrope_engine() + requests = [ + self._mrope_generation(engine, kv_cache_manager, resource_manager, 4, 1, 0, None), + self._mrope_generation(engine, kv_cache_manager, resource_manager, 6, 2, 1, None), + ] + + model_kwargs = self._forward_generation( + engine, kv_cache_manager, resource_manager, requests + ) + + # An all-zero delta vector is identical to passing no deltas at all. + self.assertNotIn("mrope_delta_read_seq_slots", model_kwargs) + self.assertNotIn("mrope_delta_write_seq_slots", model_kwargs) + # MRoPE models keep the (3,1,N) layout even with no mrope work. + expected = torch.tensor([[[3, 5]]] * 3, dtype=torch.int32, device="cuda") + torch.testing.assert_close(model_kwargs["position_ids"], expected, atol=0, rtol=0) + + def test_promoted_mrope_context_uses_decode_state_contract(self) -> None: + engine, kv_cache_manager, resource_manager, graph_runner = self._make_mrope_engine() + engine.model.config.vocab_size = 100 + context = self._request([11, 22, 33, 44], 1, seq_slot=0) + self._prefill_all_but_last_token(engine, kv_cache_manager, resource_manager, context) + mrope_delta = torch.tensor([[10]], dtype=torch.int32) + context.py_mrope_position_delta = mrope_delta + context.py_mrope_delta_cache_slot = context.py_seq_slot + context.py_multimodal_data = { + "mrope_config": { + "mrope_position_deltas": mrope_delta, + }, + "multimodal_embedding": torch.ones((1, 1), dtype=torch.float16), + } + batch = ScheduledRequests() + batch.context_requests_last_chunk = [context] + graph_runner.key = KeyType(batch_size=1, draft_len=0) + + self._forward(engine, kv_cache_manager, resource_manager, batch) + + self.assertEqual( + graph_runner.maybe_get_cuda_graph.call_args.kwargs["promoted_context_request_ids"], + frozenset({context.py_request_id}), + ) + replayed = graph_runner.replayed_inputs[0] + self.assertEqual(replayed["input_ids"].tolist(), [44]) + expected_positions = torch.full((3, 1, 1), 13, dtype=torch.int32, device="cuda") + torch.testing.assert_close(replayed["position_ids"], expected_positions, atol=0, rtol=0) + self.assertEqual(replayed["mrope_delta_read_seq_slots"].cpu().tolist(), [0]) + self.assertNotIn("mrope_delta_write_seq_slots", replayed) + attn_metadata = engine.attn_metadata + self.assertEqual(attn_metadata.kv_cache_params.num_cached_tokens_per_seq, [3]) + self.assertEqual(attn_metadata.num_contexts, 0) + + +class ModelEngineTestCase(unittest.TestCase): + def test_breakable_rejects_multimodal_models(self) -> None: + engine = object.__new__(ModelEngine) + engine.model = DummyLegacyMultimodalIndexModel() + engine.input_processor = None + engine.llm_args = SimpleNamespace( + prefill_cuda_graph_backend=PrefillCudaGraphBackend.BREAKABLE, disable_mm_encoder=False + ) + + self.assertTrue(is_multimodal(engine.model, engine.input_processor)) + with self.assertRaisesRegex(ValueError, "multimodal models"): + engine._validate_breakable_cuda_graph_compatibility() + + def test_breakable_allows_text_decoder_with_multimodal_processor(self) -> None: + engine = object.__new__(ModelEngine) + engine.model = Mock(spec=DecoderModelForCausalLM) + engine.input_processor = Mock(spec=BaseMultimodalInputProcessor) + engine.llm_args = SimpleNamespace( + prefill_cuda_graph_backend=PrefillCudaGraphBackend.BREAKABLE, disable_mm_encoder=False + ) + + self.assertTrue(is_multimodal(engine.model, engine.input_processor)) + engine._validate_breakable_cuda_graph_compatibility() + + def test_breakable_allows_multimodal_wrapper_in_text_only_mode(self) -> None: + engine = object.__new__(ModelEngine) + engine.model = DummyLegacyMultimodalIndexModel() + engine.model.llm = Mock(spec=DecoderModelForCausalLM) + engine.model.mm_encoder = None + engine.input_processor = Mock(spec=BaseMultimodalInputProcessor) + engine.llm_args = SimpleNamespace( + prefill_cuda_graph_backend=PrefillCudaGraphBackend.BREAKABLE, disable_mm_encoder=True + ) + + self.assertTrue(is_multimodal(engine.model, engine.input_processor)) + engine._validate_breakable_cuda_graph_compatibility() + + def test_restore_after_refit_never_touches_the_kv_cache(self) -> None: + """The re-wrap hook must not run the warm-up KV-cache cleanup. + + ``warmup_with_kv_cache_cleanup`` ends with + ``check_invalid_values_in_kv_cache(fill_with_zero=True)``, which + zero-fills every attention KV pool whether or not it found a NaN. That + is right after a capture warm-up that ran forwards on placeholder + inputs; restore never runs a forward, and under an in-flight weight + update the zero-fill would wipe the KV of the requests kept across the + refit. So the fully decorated hook must never even look up the KV cache + manager -- with torch.compile on or off. + """ + # Finalize earlier engines before observing this hook's GC calls. + gc.collect() + for compile_enabled in (True, False): + with self.subTest(compile_enabled=compile_enabled): + resource_manager = Mock() + model_engine = SimpleNamespace( + warmup_scope=nullcontext, + _torch_compile_backend=Mock() if compile_enabled else None, + llm_args=SimpleNamespace( + torch_compile_config=SimpleNamespace(enable_fullgraph=False) + ), + _apply_torch_compile=Mock(), + ) + with patch("gc.collect") as collect, patch("torch.compiler.reset") as reset: + ModelEngine.restore_compiled_model_after_refit(model_engine, resource_manager) + resource_manager.get_resource_manager.assert_not_called() + collect.assert_not_called() + reset.assert_not_called() + self.assertEqual( + model_engine._apply_torch_compile.call_count, 1 if compile_enabled else 0 + ) + + def test_decoder_buffers_are_allocated_only_for_decoder_runner(self) -> None: + llm_args = TorchLlmArgs( + model="dummy", + max_batch_size=4, + max_num_tokens=64, + max_beam_width=2, + cuda_graph_config=CudaGraphConfig(batch_sizes=[1, 2, 4]), + ) + for is_generation in (True, False): + with self.subTest(is_generation=is_generation): + engine = DummyModelEngine(llm_args, torch.half, is_generation=is_generation) + + runner = engine._runner + if is_generation: + self.assertIsInstance(runner, DecoderRunner) + self.assertIsNotNone(engine.cuda_graph_runner) + self.assertEqual( + tuple(runner._buffers.cache_indirection.shape), (4, 2, engine.max_seq_len) + ) + else: + self.assertIsInstance(runner, PoolingRunner) + self.assertIsNone(engine.cuda_graph_runner) + self.assertEqual(isinstance(runner._buffers, DecoderInputBuffers), is_generation) + + def test_no_kv_cache_runners_reject_speculative_config(self) -> None: + for mm_encoder_only, runner_name in ( + (False, "PoolingRunner"), + (True, "MultimodalEncoderRunner"), + ): + with self.subTest(runner=runner_name): + llm_args = TorchLlmArgs( + model="dummy", + max_batch_size=4, + max_num_tokens=64, + mm_encoder_only=mm_encoder_only, + ) + with self.assertRaisesRegex( + ValueError, f"{runner_name} does not support speculative decoding" + ): + DummyModelEngine( + llm_args, + torch.half, + spec_config=SADecodingConfig(max_draft_len=1), + is_generation=False, + ) + + @pytest.mark.cpu_only + def test_runner_engine_shared_entries_skip_decoder_state(self) -> None: + runner = Mock(spec=NoKVCacheRunner) + engine = object.__new__(ModelEngine) + engine._runner = runner + engine._torch_compile_backend = None + + engine.wait_for_input_copy() + engine.release_executor_resources(release_model_owned_graphs=False) + engine.init_cuda_graph_lora_manager(Mock()) + + runner.wait_for_input_copy.assert_called_once_with() + runner.release_graphs.assert_called_once_with(release_nccl_window_owners=True) + + def test_decoder_engine_shared_entries_use_decoder_state(self) -> None: + engine = object.__new__(ModelEngine) + decoder = object.__new__(DecoderRunner) + engine._runner = decoder + engine._torch_compile_backend = None + decoder._inputs = object.__new__(DecoderInputPreparer) + decoder._inputs._prepare_inputs_event = Mock() + decoder._graphs = DecoderCudaGraphs( + cuda_graph_runner=Mock(), + breakable_cuda_graph_runner=Mock(), + config=None, + mapping=Mapping(), + model=None, + sparse_attention_config=None, + metrics={}, + dummy_requests=None, + agreement=None, + metadata=None, + get_runtime_tokens_per_gen_step=None, + dynamic_draft_len_mapping=None, + lora=None, + forward_warmup=None, + torch_compile_bypass=nullcontext, + context_promotion_supported=True, + ) + decoder._metadata = SimpleNamespace(attn_metadata=Mock()) + + engine.wait_for_input_copy() + engine.release_executor_resources(release_model_owned_graphs=False) + + decoder._inputs._prepare_inputs_event.synchronize.assert_called_once_with() + decoder.cuda_graph_runner.clear.assert_called_once_with(release_nccl_window_owners=True) + decoder._graphs.breakable_cuda_graph_runner.clear.assert_called_once_with( + release_nccl_window_owners=True + ) + self.assertIsNone(decoder.attn_metadata) + + def test_position_id_preparation(self) -> None: + (model_engine, kv_cache_manager) = self.enterContext(create_model_engine_and_kvcache()) + resource_manager = ResourceManager({ResourceManagerType.KV_CACHE_MANAGER: kv_cache_manager}) + + prompt_len = 256 + requests = [_create_request(prompt_len, 0)] + + # Prefill run + batch = ScheduledRequests() + batch.context_requests_last_chunk = requests + _prepare_kv_cache_resources(kv_cache_manager, batch) + model_engine.forward(batch, resource_manager) + + expected_prefill_pos_ids = torch.arange( + 0, prompt_len, dtype=torch.int32, device="cuda" + ).unsqueeze(0) + torch.testing.assert_close( + model_engine.model.recorded_position_ids, expected_prefill_pos_ids, atol=0, rtol=0 + ) + + # Simulate decoding one token after prefill + requests[-1].add_new_token(42, 0) + + # Generation run + batch = ScheduledRequests() + batch.generation_requests = requests + _prepare_kv_cache_resources(kv_cache_manager, batch) + + model_engine.forward(batch, resource_manager) + expected_gen_pos_id = torch.tensor( + [prompt_len], dtype=torch.int32, device="cuda" + ).unsqueeze(0) + torch.testing.assert_close( + model_engine.model.recorded_position_ids, expected_gen_pos_id, atol=0, rtol=0 + ) + + def test_warmup(self) -> None: + (model_engine, kv_cache_manager) = self.enterContext(create_model_engine_and_kvcache()) + resource_manager = ResourceManager({ResourceManagerType.KV_CACHE_MANAGER: kv_cache_manager}) + + # Test with a huge batch size. The warmup run should bail out of + # warmup instead of crashing (there's not enough KV cache space for this). + model_engine._runner._config.cuda_graph_batch_sizes.append(1000000000) + + num_free_before = kv_cache_manager.get_num_free_blocks() + model_engine.warmup(resource_manager) + + # Warmup pre-allocates the CUDA graph padding dummy for each captured + # draft length (only draft_len 0 here, no speculation), which + # intentionally keeps holding its KV blocks so padding cannot fall + # back to eager once the cache saturates. + padding_dummies = model_engine.cuda_graph_runner.padding_dummy_requests + self.assertEqual([0], sorted(padding_dummies.keys())) + + # Make sure we don't leak any blocks beyond that dummy: freeing it + # must restore the exact pre-warmup free-block count. + kv_cache_manager.free_resources(padding_dummies[0]) + self.assertEqual(num_free_before, kv_cache_manager.get_num_free_blocks()) + + def test_warmup_skips_padding_dummy_when_padding_impossible(self) -> None: + # With max_batch_size=1 and a graph for batch size 1, every batch + # already matches a graph size, so no padded batch can ever occur and + # warmup must not retain a padding dummy (it would permanently hold + # KV blocks and spec/hybrid resources that are never used). + llm_args = TorchLlmArgs( + model="dummy", + max_batch_size=1, + max_num_tokens=258, + cuda_graph_config=CudaGraphConfig(enable_padding=True, batch_sizes=[1, 2, 4]), + ) + (model_engine, kv_cache_manager) = self.enterContext( + create_model_engine_and_kvcache(llm_args) + ) + resource_manager = ResourceManager({ResourceManagerType.KV_CACHE_MANAGER: kv_cache_manager}) + + num_free_before = kv_cache_manager.get_num_free_blocks() + model_engine.warmup(resource_manager) + + self.assertEqual({}, model_engine.cuda_graph_runner.padding_dummy_requests) + self.assertEqual(num_free_before, kv_cache_manager.get_num_free_blocks()) + + def test_warmup_skips_padding_dummy_during_estimation(self) -> None: + # The estimation-phase KV cache is sized with no headroom for retained + # dummies; holding blocks there can leave the estimation requests + # unschedulable. Warmup must skip the preallocation for managers + # created for KV cache capacity estimation. + (model_engine, kv_cache_manager) = self.enterContext(create_model_engine_and_kvcache()) + kv_cache_manager.is_estimating_kv_cache = True + resource_manager = ResourceManager({ResourceManagerType.KV_CACHE_MANAGER: kv_cache_manager}) + + num_free_before = kv_cache_manager.get_num_free_blocks() + model_engine.warmup(resource_manager) + + self.assertEqual({}, model_engine.cuda_graph_runner.padding_dummy_requests) + self.assertEqual(num_free_before, kv_cache_manager.get_num_free_blocks()) + + def test_cuda_graph_padding_filters_huge_batch_size(self) -> None: + llm_args = TorchLlmArgs( + model="dummy", + cuda_graph_config=CudaGraphConfig( + enable_padding=True, batch_sizes=[1, 2, 3, 1000000000000000000000000] + ), + ) + model_engine = DummyModelEngine(llm_args, torch.half) + + self.assertEqual(model_engine.cuda_graph_batch_sizes, [1, 2, 3, model_engine.max_seq_len]) + self.assertTrue(model_engine.cuda_graph_padding_enabled) + + def test_kv_cache_manager_with_execution_stream(self) -> None: + """Test that KVCacheManagerV2 uses the provided execution_stream.""" + # Create a dedicated execution stream + execution_stream = torch.cuda.Stream() + + (model_engine, kv_cache_manager) = self.enterContext( + create_model_engine_and_kvcache(execution_stream=execution_stream) + ) + + # Verify the KVCacheManagerV2 uses the provided execution stream + self.assertEqual( + kv_cache_manager._stream.cuda_stream, + execution_stream.cuda_stream, + "KVCacheManagerV2 should use the provided execution_stream", + ) + + resource_manager = ResourceManager({ResourceManagerType.KV_CACHE_MANAGER: kv_cache_manager}) + + prompt_len = 32 + requests = [_create_request(prompt_len, 0)] + + batch = ScheduledRequests() + batch.context_requests_last_chunk = requests + _prepare_kv_cache_resources(kv_cache_manager, batch) + with torch.cuda.stream(execution_stream): + model_engine.forward(batch, resource_manager) + + # Verify the stream is still the same after forward pass + self.assertEqual( + kv_cache_manager._stream.cuda_stream, + execution_stream.cuda_stream, + "KVCacheManagerV2 should still use the provided execution_stream after forward", + ) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("model_kind", ["default", "m3", "m3_vl", "gdn", "mamba"]) +@pytest.mark.parametrize("compile_enabled,piecewise", [(False, False), (True, False), (True, True)]) +@pytest.mark.parametrize("compile_only_piecewise_graphs", [False, True]) +def test_pcg_fx_fallback_policy_is_model_specific( + monkeypatch: pytest.MonkeyPatch, + model_kind: str, + compile_enabled: bool, + piecewise: bool, + compile_only_piecewise_graphs: bool, +) -> None: + """Exercise the real constructor's routing gate without loading weights or CUDA.""" + from tensorrt_llm._torch.models.modeling_minimaxm3 import ( + MiniMaxM3ForCausalLM, + MiniMaxM3VLForConditionalGeneration, + ) + from tensorrt_llm._torch.models.modeling_nemotron_h import NemotronHForCausalLM + from tensorrt_llm._torch.models.modeling_qwen3_next import Qwen3NextForCausalLM + from tensorrt_llm._torch.pyexecutor import _util + + model_cls = { + "default": DecoderModelForCausalLM, + "m3": MiniMaxM3ForCausalLM, + "m3_vl": MiniMaxM3VLForConditionalGeneration, + "gdn": Qwen3NextForCausalLM, + "mamba": NemotronHForCausalLM, + }[model_kind] + assert model_cls.use_fx_for_pcg_fallback is (model_kind not in ("m3", "m3_vl")) + model = model_cls.__new__(model_cls) + torch.nn.Module.__init__(model) + mapping = Mapping() + model.model_config = SimpleNamespace( + pretrained_config=SimpleNamespace(torch_dtype=torch.float32, hidden_size=4), + mapping=mapping, + sparse_attention_config=None, + ) + eager = torch.nn.Linear(4, 4) + model.model = eager + compile_config = SimpleNamespace( + enable_fullgraph=True, + compile_only_piecewise_graphs=compile_only_piecewise_graphs, + enable_inductor=False, + enable_userbuffers=False, + max_num_streams=1, + ) + llm_args = SimpleNamespace( + encode_only=False, + mm_encoder_only=False, + get_runtime_sizes=lambda: (1, 16, 64, 4), + encoder_max_batch_size=None, + encoder_max_num_tokens=None, + enable_in_graph_sampling=False, + enable_return_routed_experts=False, + multimodal_config=MultimodalConfig(), + checkpoint_format="HF", + trust_remote_code=False, + disable_overlap_scheduler=True, + kv_cache_config=KvCacheConfig(), + sleep_config=None, + enable_layerwise_nvtx_marker=False, + cuda_graph_config=None, + torch_compile_config=compile_config if compile_enabled else None, + prefill_cuda_graph_backend=PrefillCudaGraphBackend.PIECEWISE if piecewise else None, + prefill_capture_num_tokens=[4], + allreduce_strategy="AUTO", + attn_backend="TRTLLM", + sparse_attention_config=None, + ) + for name in ( + "should_enable_overlap_headroom", + "resolved_kv_cache_manager_is_v2", + ): + monkeypatch.setattr(_util, name, Mock(return_value=False)) + monkeypatch.setattr(_util, "compute_max_num_sequences", Mock(return_value=4)) + for name in ( + "_configure_deep_gemm_pdl", + "create_input_processor", + "setup_mm_encoder_attn_metadata", + ): + monkeypatch.setattr(model_engine_module, name, Mock()) + monkeypatch.setattr( + model_engine_module, "resolve_mrope_position_deltas_cache", lambda model: None + ) + monkeypatch.setattr( + model_engine_module, "is_hybrid_linear", lambda config: model_kind in ("gdn", "mamba") + ) + monkeypatch.setattr( + model_engine_module.MultimodalItemScheduler, "maybe_create", Mock(return_value=None) + ) + monkeypatch.setattr(ModelEngine, "_validate_breakable_cuda_graph_compatibility", Mock()) + monkeypatch.setattr(ModelEngine, "_init_model_capacity", Mock()) + monkeypatch.setattr(ModelEngine, "__del__", lambda self: None) + backend_factory = Mock() + backend_factory.Streams = list + monkeypatch.setattr(model_engine_module, "Backend", backend_factory) + compiled = torch.nn.Identity() + compile_model = Mock(return_value=compiled) + monkeypatch.setattr(torch, "compile", compile_model) + monkeypatch.setattr(torch._dynamo.config, "cache_size_limit", 16) + # Stop after the complete compilation block, before runtime cache allocation. + monkeypatch.setattr( + model_engine_module, + "get_attention_backend", + Mock(side_effect=RuntimeError("compile setup complete")), + ) + engine = ModelEngine.__new__(ModelEngine) + with torch_compiling(False), pytest.raises(RuntimeError, match="compile setup complete"): + ModelEngine.__init__( + engine, + model_path="dummy", + mapping=mapping, + model=model, + llm_args=llm_args, + checkpoint_loader=Mock(), + ) + + expected_prefill_only = ( + compile_enabled + and piecewise + and (compile_only_piecewise_graphs or model_kind in ("m3", "m3_vl")) + ) + assert engine._torch_compile_prefill_only is expected_prefill_only + if not compile_enabled: + compile_model.assert_not_called() + assert model.model is eager + else: + compile_model.assert_called_once_with( + eager, backend=engine._torch_compile_backend, fullgraph=True + ) + assert backend_factory.call_args.args[0] is False # Inductor stays disabled. + if expected_prefill_only: + assert isinstance(model.model, PrefillCompiledModel) + assert model.model.eager_model is eager + assert model.model.compiled_model is compiled + else: + assert model.model is compiled + + +class TestWarmupCleanup(unittest.TestCase): + @pytest.mark.cpu_only + def test_no_kv_cache_warmup_rejects_allocated_kv_cache(self) -> None: + """Verify a no-KV-cache runner rejects an allocated KV cache before warmup.""" + model_engine = object.__new__(ModelEngine) + model_engine._metrics = {} + model_engine.model = SimpleNamespace(model_config=SimpleNamespace(is_encoder_decoder=False)) + model_engine.moe_load_balancer = None + model_engine._warmup_depth = 0 + model_engine._runner = Mock(spec=NoKVCacheRunner) + runner = model_engine._runner + runner._validate_resources = NoKVCacheRunner._validate_resources.__get__(runner) + runner.warmup.side_effect = NoKVCacheRunner.warmup.__get__(runner) + resource_manager = Mock() + resource_manager.get_resource_manager.return_value = object() + + with self.assertRaisesRegex( + AssertionError, + "no-KV-cache runner was initialized, but a KV cache manager was allocated", + ): + model_engine.warmup(resource_manager) + + model_engine._runner.warmup.assert_called_once_with(resource_manager) + resource_manager.get_resource_manager.assert_called_once_with( + ResourceManagerType.KV_CACHE_MANAGER + ) + + @pytest.mark.cpu_only + def test_warmup_total_timing_and_kv_cache_cleanup(self) -> None: + model_engine = object.__new__(ModelEngine) + model_engine._metrics = {} + model_engine._warmup_depth = 0 + model_engine.moe_load_balancer = None + model_engine._runner = Mock(spec=DecoderRunner) + kv_cache_manager = Mock() + kv_cache_manager.check_invalid_values_in_kv_cache.return_value = False + resource_manager = Mock() + resource_manager.get_resource_manager.return_value = kv_cache_manager + with ( + patch.object(model_engine, "_warmup_impl") as warmup_impl, + patch("time.perf_counter", side_effect=[0.0, 2.0, 5.0, 7.0]), + ): + model_engine.warmup(resource_manager) + warmup_impl.assert_called_once_with(resource_manager) + resource_manager.get_resource_manager.assert_called_once_with( + ResourceManagerType.KV_CACHE_MANAGER + ) + kv_cache_manager.check_invalid_values_in_kv_cache.assert_called_once_with( + fill_with_zero=True + ) + self.assertEqual( + model_engine.metrics, {"kv_cache_cleanup_seconds": 3.0, "total_warmup_seconds": 7.0} + ) + + def test_encoder_decoder_encoder_warmup_delegates_runner_lifecycle(self) -> None: + model_engine = object.__new__(ModelEngine) + model_engine.model = SimpleNamespace(model_config=SimpleNamespace(is_encoder_decoder=True)) + model_engine.moe_load_balancer = None + model_engine._warmup_depth = 0 + model_engine._runner = Mock(spec=EncoderDecoderRunner) + warmup_depths = [] + model_engine._runner.warmup_encoder.side_effect = lambda _: warmup_depths.append( + model_engine._warmup_depth + ) + resource_manager = object() + + model_engine.warmup_encoder(resource_manager) + + self.assertEqual( + model_engine._runner.method_calls, + [call.warmup_encoder(resource_manager)], + ) + self.assertEqual(warmup_depths, [1]) + self.assertEqual(model_engine._warmup_depth, 0) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + ("cuda_graph_config", "expected_encode_config"), + [ + (None, False), + (CudaGraphConfig(), False), + (EncodeCudaGraphConfig(batch_sizes=[1], num_tokens=[16], seq_lens=[16]), True), + ], +) +def test_encoder_runner_graph_config_comes_only_from_cuda_graph_config( + cuda_graph_config: CudaGraphConfig | None, + expected_encode_config: bool, +) -> None: + """Encode-only takes its buckets from `cuda_graph_config`, never from llm_args.""" + engine = object.__new__(ModelEngine) + engine._model_caller = Mock() + engine.model = object() + engine.mapping = object() + engine.cuda_graph_config = cuda_graph_config + engine.llm_args = SimpleNamespace( + encoder_cuda_graph_config=EncodeCudaGraphConfig( + batch_sizes=[8], + num_tokens=[64], + seq_lens=[64], + ), + enable_autotuner=False, + ) + engine.batch_size = 4 + engine.max_num_tokens = 16 + engine.max_seq_len = 16 + engine.max_beam_width = 1 + engine.without_logits = False + engine.attn_backend = _AttentionBackend + engine.attn_runtime_features = AttentionRuntimeFeatures() + engine.dist = object() + engine.moe_load_balancer = None + runner_config = object() + expected_runner = object() + runner_type = Mock(return_value=expected_runner) + + with patch( + "tensorrt_llm._torch.pyexecutor.model_engine.EncoderRunnerConfig.create", + return_value=runner_config, + ) as create_config: + actual = engine._initialize_encoder_runner(runner_type) + + assert actual is expected_runner + resolved = create_config.call_args.kwargs["graph_config"] + assert resolved is (cuda_graph_config if expected_encode_config else None) + runner_type.assert_called_once_with( + engine.model, + runner_config, + mapping=engine.mapping, + dist=engine.dist, + moe_load_balancer=engine.moe_load_balancer, + model_caller=engine._model_caller, + attention_runtime_features=engine.attn_runtime_features, + ) + + +@pytest.mark.cpu_only +def test_model_engine_forward_encoder_delegates_scheduled_encoder_batch() -> None: + runner = Mock(spec=EncoderDecoderRunner) + hidden_states = object() + runner.forward_encoder.return_value = { + "encoder_hidden_states": hidden_states, + "encoder_seq_lens": [2, 3], + } + engine = object.__new__(ModelEngine) + engine._model_caller = Mock() + engine.model = SimpleNamespace( + model_config=SimpleNamespace(is_encoder_decoder=True), + ) + engine._runner = runner + engine._warmup_depth = 0 + requests = [object(), object()] + resource_manager = object() + + outputs = engine.forward_encoder(requests, resource_manager) + + assert outputs == (hidden_states, [2, 3]) + inputs = runner.forward_encoder.call_args.args[0] + assert inputs.batch.encoder_requests == requests + runner.forward_encoder.assert_called_once_with( + inputs, + resource_manager=resource_manager, + ) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("enable_spec_decode", [False, True]) +@pytest.mark.parametrize("effective_length", [None, 0, 3]) +@pytest.mark.parametrize("is_dummy", [False, True]) +def test_engine_passes_call_state_and_returns_length_update( + effective_length: int | None, + is_dummy: bool, + enable_spec_decode: bool, +) -> None: + runner = Mock(spec=NoKVCacheRunner) + outputs = {"logits": object()} + if effective_length is not None: + outputs["runtime_draft_len"] = effective_length + runner.forward.return_value = outputs + engine, resources = _model_engine_with_runner(runner, kv_cache_manager=None) + engine._warmup_depth = int(is_dummy) + batch = ScheduledRequests() + batch.context_requests_last_chunk = [SimpleNamespace(py_return_context_logits=True)] + new_tokens = SimpleNamespace(new_tokens=torch.tensor([17])) + cache_indirection = torch.tensor([0], dtype=torch.int32) + runtime_draft_len = 2 if enable_spec_decode else 0 + assert ( + engine.forward( + batch, + resources, + new_tokens, + cache_indirection, + enable_spec_decode=enable_spec_decode, + runtime_draft_len=runtime_draft_len, + ) + is outputs + ) + inputs = runner.forward.call_args.args[0] + assert inputs.batch is batch + assert inputs.new_tensors_device is new_tokens + assert inputs.cache_indirection_buffer is cache_indirection + assert inputs.gather_context_logits + assert inputs.enable_spec_decode is enable_spec_decode + assert inputs.runtime_draft_len == runtime_draft_len + assert runner.forward.call_args.kwargs["is_dummy"] is is_dummy + assert outputs.get("runtime_draft_len") == effective_length + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("raw_output", [torch.tensor([1.0]), None], ids=["tensor", "none"]) +@pytest.mark.parametrize("is_dummy", [False, True]) +def test_engine_forward_preserves_raw_decoder_outputs( + raw_output: torch.Tensor | None, is_dummy: bool +) -> None: + engine, resources = _model_engine_with_runner(None, kv_cache_manager=object()) + engine._warmup_depth = int(is_dummy) + engine._runner._forward_decoder = Mock(return_value=raw_output) + + assert engine.forward(ScheduledRequests(), resources, runtime_draft_len=2) is raw_output + assert engine._runner._forward_decoder.call_args.args[0].runtime_draft_len == 2 + assert engine._runner._forward_decoder.call_args.kwargs == {"is_dummy": is_dummy} + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("compile_enabled", [False, True]) +@pytest.mark.parametrize("prefill_only", [False, True]) +def test_refit_updates_compiled_weights_without_invalidating_graphs( + compile_enabled: bool, + prefill_only: bool, +) -> None: + model = DecoderModelForCausalLM.__new__(DecoderModelForCausalLM) + torch.nn.Module.__init__(model) + eager = model.model = torch.nn.Linear(2, 2, bias=False) + storage = eager.weight.data_ptr() + engine = object.__new__(ModelEngine) + engine._cleanup_done = True + engine.model = model + engine._torch_compile_backend = "eager" if compile_enabled else None + engine._warmup_depth = 0 + engine.moe_load_balancer = None + engine.llm_args = SimpleNamespace( + prefill_cuda_graph_backend=PrefillCudaGraphBackend.PIECEWISE + if prefill_only + else PrefillCudaGraphBackend.DISABLED, + torch_compile_config=SimpleNamespace( + compile_only_piecewise_graphs=prefill_only, enable_fullgraph=False + ), + ) + graphs = {KeyType(batch_size=1, draft_len=0): object()} + padding = {0: object()} + engine._runner = Mock(spec=DecoderRunner) + engine._runner.cuda_graph_runner = SimpleNamespace( + graphs=graphs.copy(), padding_dummy_requests=padding.copy() + ) + resources = Mock() + inputs = torch.tensor([[1.0, 2.0]]) + with ( + patch("torch.cuda.synchronize") as synchronize, + patch( + "tensorrt_llm._torch.pyexecutor.engine.runners.decoder.runner.get_per_request_prefill_cuda_graph_flag", + return_value=True, + ), + ): + if compile_enabled: + engine._apply_torch_compile("eager", False) + for value, unwrap, restore in ( + ( + 2.0, + engine.unwrap_compiled_model_for_refit, + engine.restore_compiled_model_after_refit, + ), + ( + 3.0, + engine.release_piecewise_cuda_graphs_for_refit, + engine.recapture_piecewise_cuda_graphs_after_refit, + ), + ): + synchronize.reset_mock() + unwrap() + if compile_enabled: + synchronize.assert_called_once_with() + else: + synchronize.assert_not_called() + replacement = torch.full_like(eager.weight, value) + model.load_state_dict({"model.weight": replacement}) + restore(resources) + torch.testing.assert_close(model.model(inputs), inputs @ replacement.t()) + assert eager.weight.data_ptr() == storage + assert engine.cuda_graph_runner.graphs == graphs + assert engine.cuda_graph_runner.padding_dummy_requests == padding + resources.get_resource_manager.assert_not_called() + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("enabled", [False, True]) +@pytest.mark.parametrize("architectures", [None, ["Decoder", "Wrapper"]]) +def test_layerwise_nvtx_flag_controls_real_forward_markers( + monkeypatch: pytest.MonkeyPatch, + enabled: bool, + architectures: list[str] | None, +) -> None: + from tensorrt_llm._torch.pyexecutor import _util + + model = torch.nn.Sequential(torch.nn.Linear(2, 2, bias=False), torch.nn.ReLU()) + model.model_config = DummyModel(torch.float32).model_config + model.config = model.model_config.pretrained_config + model.config.architectures = architectures + with torch.no_grad(): + model[0].weight.fill_(1) + engine = object.__new__(ModelEngine) + engine._cleanup_done = True + engine.model = model + engine.mapping = Mapping() + engine.dist = None + engine.enable_in_graph_sampling = False + engine.max_beam_width = 1 + engine.batch_size = 4 + engine.is_spec_decode = False + engine.input_processor = None + engine.encoder_max_num_tokens = 16 + engine.llm_args = TorchLlmArgs( + model="dummy", disable_overlap_scheduler=True, enable_layerwise_nvtx_marker=enabled + ) + engine._init_model_capacity = Mock() + monkeypatch.setattr(_util, "should_enable_overlap_headroom", Mock(return_value=False)) + monkeypatch.setattr(_util, "resolved_kv_cache_manager_is_v2", Mock(return_value=False)) + monkeypatch.setattr(_util, "compute_max_num_sequences", Mock(return_value=4)) + events = [] + monkeypatch.setattr( + torch.cuda.nvtx, + "range_push", + lambda label: events.append(("push", literal_eval(label)["Module"])), + ) + monkeypatch.setattr(torch.cuda.nvtx, "range_pop", lambda: events.append(("pop", None))) + engine._init_from_model() + torch.testing.assert_close(model(torch.tensor([[1.0, 2.0]])), torch.tensor([[3.0, 3.0]])) + prefix = "|".join(architectures) if architectures else "Model" + assert events == ( + [ + ("push", prefix), + ("push", prefix + ".0"), + ("pop", None), + ("push", prefix + ".1"), + ("pop", None), + ("pop", None), + ] + if enabled + else [] + ) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("fails", [False, True]) +def test_warmup_restores_serving_state_after_success_or_failure(fails: bool) -> None: + runner = Mock(spec=NoKVCacheRunner) + engine, resources = _model_engine_with_runner(runner, kv_cache_manager=None) + engine._metrics = {} + balancer = engine.moe_load_balancer = Mock() + batch = ScheduledRequests() + + def warmup(resource_manager: ResourceManager) -> None: + assert resource_manager is resources + engine.forward(batch, resources) + if fails: + raise RuntimeError("warmup failed") + + runner.warmup.side_effect = warmup + with patch.object(model_engine_module, "set_moe_comm_warmup") as communication_mode: + with engine.warmup_scope(): + if fails: + with pytest.raises(RuntimeError, match="warmup failed"): + engine.warmup(resources) + else: + engine.warmup(resources) + engine.forward(batch, resources) + assert balancer.set_iter_info.call_args_list == [ + call(enable_statistic=False, enable_update_weights=False) + ] + engine.forward(batch, resources) + assert [entry.kwargs["is_dummy"] for entry in runner.forward.call_args_list] == [ + True, + True, + False, + ] + assert balancer.set_iter_info.call_args_list == [ + call(enable_statistic=False, enable_update_weights=False), + call(enable_statistic=True, enable_update_weights=True), + ] + assert communication_mode.call_args_list == [call(True), call(False)] diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py deleted file mode 100644 index 4be8c6a539a2..000000000000 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ /dev/null @@ -1,4224 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - -import unittest -import weakref -from contextlib import nullcontext -from dataclasses import dataclass, replace -from types import SimpleNamespace -from unittest.mock import Mock, call, patch - -import pytest -import torch - -import tensorrt_llm -from tensorrt_llm._torch.model_config import ModelConfig -from tensorrt_llm._torch.models.modeling_multimodal_encoder import \ - MultimodalEncoderMixin -from tensorrt_llm._torch.models.modeling_multimodal_mixin import ( - MultimodalModelMixin, _build_request_multimodal_input) -from tensorrt_llm._torch.models.modeling_utils import DecoderModelForCausalLM -from tensorrt_llm._torch.pyexecutor.connectors.kv_cache_connector import \ - KvCacheConnectorWorker -from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import ( - CUDAGraphRunner, KeyType, SampleType, _restore_spec_decode_capture_state, - _save_spec_decode_capture_state) -from tensorrt_llm._torch.pyexecutor.engine.input_buffers import InputBuffers -from tensorrt_llm._torch.pyexecutor.engine.model_call import ModelCaller -from tensorrt_llm._torch.pyexecutor.engine.multimodal import ( - is_multimodal, setup_mm_encoder_attn_metadata) -from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import DecoderRunner -from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.graphs import ( - DecoderCudaGraphs, _make_single_token_context_graph_batch) -from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.inputs import ( - DecoderInputBuffers, DecoderInputPreparer, HostInputs, - _get_context_prompt_lookahead_token, _zero_fill_row_spans) -from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.metadata import \ - DecoderMetadata -from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.runner import \ - PrefillCompiledModel -from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.warmup import \ - DecoderWarmup -from tensorrt_llm._torch.pyexecutor.engine.runners.encoder_decoder import \ - EncoderDecoderRunner -from tensorrt_llm._torch.pyexecutor.engine.runners.interface import \ - PreparedInputs -from tensorrt_llm._torch.pyexecutor.engine.runners.pooling import PoolingRunner -from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import \ - KVCacheManagerV2 -from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest -from tensorrt_llm._torch.pyexecutor.model_engine import ModelEngine -from tensorrt_llm.llmapi.llm_args import (DecodingBaseConfig, - PrefillCudaGraphBackend, - SeqLenAwareSparseAttentionConfig, - TorchLlmArgs) - -# isort: off -from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManager, ResourceManagerType -# isort: on -from utils.util import skip_pre_blackwell_unittest, skip_ray - -from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata -from tensorrt_llm._torch.attention.backends.trtllm import \ - TrtllmAttentionMetadata -from tensorrt_llm._torch.attention.backends.vanilla import VanillaAttention -from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests -from tensorrt_llm._torch.speculative.interface import \ - INVALID_PROMPT_LOOKAHEAD_TOKEN -from tensorrt_llm._torch.speculative.spec_sampler_base import \ - SampleStateTensorsSpec -from tensorrt_llm._torch.utils import set_per_request_prefill_cuda_graph_flag -from tensorrt_llm.inputs.multimodal import ( - MULTIMODAL_ENCODER_ITEM_METADATA_KEY, MultimodalItemMetadata) -from tensorrt_llm.inputs.registry import (BaseMultimodalDummyInputsBuilder, - BaseMultimodalInputProcessor) -from tensorrt_llm.llmapi import (CudaGraphConfig, MTPDecodingConfig, - NGramDecodingConfig, SADecodingConfig, - SamplingParams) -from tensorrt_llm.llmapi.llm_args import KvCacheConfig -from tensorrt_llm.mapping import CpType, Mapping - -# One-engine checks see a two-model mode, so no draft KV cache is looked up. -_TWO_MODEL_SPEC_MODE = SimpleNamespace(use_one_engine=lambda: False) - - -def _input_preparer(*, without_config: bool = False) -> DecoderInputPreparer: - return DecoderInputPreparer( - buffers=None, - config=(None if without_config else SimpleNamespace( - is_spec_decode=False, spec_config=None)), - mapping=Mapping(), - dist=None, - model=None, - get_runtime_tokens_per_gen_step=None, - iter_states={}, - lora=None, - use_beam_search=False, - steady_gen_cache_supported=True) - - -@pytest.mark.parametrize("runner_type,graphs_path", [ - (DecoderRunner, - "tensorrt_llm._torch.pyexecutor.engine.runners.decoder.runner.DecoderCudaGraphs" - ), - (EncoderDecoderRunner, - "tensorrt_llm._torch.pyexecutor.engine.runners.encoder_decoder.EncoderDecoderCudaGraphs" - ) -]) -def test_cuda_graph_factory_preserves_input_preparer(runner_type, graphs_path): - runner = Mock() - runner._inputs = object() - runner._cuda_graphs_kwargs.side_effect = lambda: DecoderRunner._cuda_graphs_kwargs( - runner) - - with patch(graphs_path) as graphs_type: - graphs = runner_type._make_cuda_graphs(runner) - - graphs_type.assert_called_once() - assert graphs_type.call_args.kwargs["inputs"] is runner._inputs - assert graphs_type.call_args.kwargs[ - "warmup_inputs"] is runner._buffers.input_ids_cuda - assert graphs is graphs_type.return_value - - -@pytest.mark.parametrize("enabled,fused", [(None, False), (False, True), - (True, False), (True, True)]) -def test_input_preparer_configured_confidence_flags(enabled, fused): - spec_config = (None if enabled is None else SimpleNamespace( - enable_confidence_scheduling=enabled, - enable_fused_confidence_scheduler=fused)) - config = SimpleNamespace(is_spec_decode=False, spec_config=spec_config) - inputs = DecoderInputPreparer(buffers=None, - config=config, - mapping=Mapping(), - dist=None, - model=None, - get_runtime_tokens_per_gen_step=None, - iter_states={}, - lora=None, - use_beam_search=False, - steady_gen_cache_supported=True) - - assert inputs._config is config - assert inputs._dspark_confidence_enabled is bool(enabled) - assert inputs._dspark_trims_submitted_tokens is bool(enabled) - assert inputs._dspark_device_windows is bool(enabled and fused) - assert inputs.previous_verify_lens == ([] if enabled else None) - assert inputs.has_previous_device_draft is False - - -def test_input_preparer_without_config_keeps_confidence_disabled(): - inputs = _input_preparer(without_config=True) - - assert inputs._config is None - assert inputs._dspark_confidence_enabled is False - assert inputs._dspark_trims_submitted_tokens is False - assert inputs._dspark_device_windows is False - assert inputs.previous_verify_lens is None - assert inputs.has_previous_device_draft is False - - -class TestCachedKvTokenLogging(unittest.TestCase): - - def setUp(self) -> None: - self.inputs = _input_preparer() - self.real = SimpleNamespace(is_cuda_graph_dummy=False) - self.padding = SimpleNamespace(is_cuda_graph_dummy=True) - - def check_counts( - self, - counts: list[int], - groups: tuple[tuple[list[SimpleNamespace], int], ...], - expected: list[int], - padding: int, - ) -> None: - self.inputs.iter_states['cached_kv_tokens'] = sum(counts) - self.inputs.record_cached_kv_tokens_per_req(counts, groups) - self.assertEqual( - self.inputs.iter_states, { - 'cached_kv_tokens': sum(counts), - 'cached_kv_tokens_per_req': expected, - 'cached_kv_tokens_cuda_graph_padding': padding, - }) - counts[:] = [999] - self.assertEqual(self.inputs.iter_states['cached_kv_tokens_per_req'], - expected) - - def test_mtp_padding_before_generation(self) -> None: - # Input packing is context, extend (including padding), gen. - self.check_counts([10, 100, 7, 200], - (([self.real], 1), ([self.real, self.padding], 1), - ([self.real], 1)), [10, 100, 200], 7) - - def test_beams_preserved_after_context_and_padding(self) -> None: - self.check_counts([10, 100, 101, 7, 7, 200, 201], - (([self.real], 1), - ([self.real, self.padding, self.real], 2)), - [10, 100, 101, 200, 201], 14) - - def test_attention_dp_dummy_is_not_graph_padding(self) -> None: - adp_dummy = SimpleNamespace(is_dummy=True, - is_attention_dp_dummy=True, - is_cuda_graph_dummy=False) - self.check_counts([10, 1, 7], - (([self.real, adp_dummy, self.padding], 1), ), - [10, 1], 7) - - def test_padding_only(self) -> None: - self.check_counts([7, 7], (([self.padding, self.padding], 1), ), [], 14) - - def test_group_row_count_must_match_counts(self) -> None: - for counts, groups in ( - ([100, 200], (([self.real], 1), )), - ([100, 101, 200], (([self.real, self.real], 2), )), - ): - with self.subTest(counts=counts, groups=groups): - with self.assertRaises(AssertionError): - self.inputs.record_cached_kv_tokens_per_req(counts, groups) - - def test_steady_batch_clears_previous_padding(self) -> None: - self.inputs.iter_states['cached_kv_tokens_cuda_graph_padding'] = 7 - self.check_counts([100, 200], (), [100, 200], 0) - - def test_empty_batch_clears_previous_values(self) -> None: - self.inputs.iter_states['cached_kv_tokens_per_req'] = [100] - self.inputs.iter_states['cached_kv_tokens_cuda_graph_padding'] = 7 - self.check_counts([], (([], 1), ), [], 0) - - -@dataclass -class Config: - torch_dtype: torch.dtype - num_key_value_heads: int = 16 - num_attention_heads: int = 16 - hidden_size: int = 256 - architectures: list[str] = None - rope_scaling: dict | None = None - - @property - def head_dim(self) -> int: - return self.hidden_size // self.num_attention_heads - - -class DummyKvCacheConnectorWorker(KvCacheConnectorWorker): - - def register_kv_caches(self, kv_cache_tensor: torch.Tensor): - pass - - -class DummyModel(torch.nn.Module): - - def __init__(self, dtype: torch.dtype, rope_scaling: dict | None = None): - super().__init__() - self.model_config = ModelConfig(pretrained_config=Config( - torch_dtype=dtype, rope_scaling=rope_scaling)) - self.recorded_position_ids = None - self.recorded_kwargs = [] - - def infer_max_seq_len(self): - return 2048 - - @property - def config(self): - return self.model_config.pretrained_config - - def forward(self, *args, **kwargs) -> torch.Tensor: - input_ids = kwargs["input_ids"] - self.recorded_position_ids = kwargs["position_ids"] - # Input buffers are reused across steps, so keep copies. - if not torch.cuda.is_current_stream_capturing(): - self.recorded_kwargs.append(_clone_tensors(kwargs)) - batch_size = input_ids.size(0) - return {"logits": torch.randn((batch_size, 10), device='cuda')} - - -class DummyLegacyMultimodalIndexModel(MultimodalModelMixin, torch.nn.Module): - - class Config: - vocab_size = 100 - - config = Config() - - @property - def mm_token_ids(self) -> torch.Tensor: - return torch.tensor([90, 91], dtype=torch.int32) - - -class DummyModelEngine(ModelEngine): - - def __init__( - self, - llm_args: TorchLlmArgs, - dtype: torch.dtype, - spec_config: DecodingBaseConfig | None = None, - is_generation: bool = True, - model: DummyModel | None = None, - ) -> None: - self.dtype = dtype - mapping = Mapping(world_size=tensorrt_llm.mpi_world_size(), - tp_size=tensorrt_llm.mpi_world_size(), - rank=tensorrt_llm.mpi_rank()) - if model is None: - model = DummyModel(self.dtype) - model.model_config.is_generation = is_generation - super().__init__(model_path="dummy", - mapping=mapping, - model=model, - llm_args=llm_args, - spec_config=spec_config) - - -def _prepare_kv_cache_resources(manager: KVCacheManagerV2, - batch: ScheduledRequests) -> None: - # These model-engine tests bypass the scheduler, which owns V2 allocation. - for request in batch.context_requests: - assert manager.prepare_context(request) - assert manager.resize_context( - request, request.prompt_len - request.context_current_position) - for request in batch.generation_requests: - assert manager.try_allocate_generation(request) - manager.prepare_resources(batch) - - -def _create_request(num_tokens, req_id: int): - sampling_params = SamplingParams() - kwargs = { - "request_id": - req_id, - "max_new_tokens": - 1, - "input_tokens": [0] * num_tokens, - "sampling_config": - tensorrt_llm.bindings.SamplingConfig( - sampling_params._get_sampling_config()), - "is_streaming": - False, - } - result = LlmRequest(**kwargs) - result.paged_kv_block_ids = [] - return result - - -def _create_request_with_tokens(tokens: list[int], req_id: int) -> LlmRequest: - sampling_params = SamplingParams() - request = LlmRequest( - request_id=req_id, - max_new_tokens=1, - input_tokens=tokens, - sampling_config=tensorrt_llm.bindings.SamplingConfig( - sampling_params._get_sampling_config()), - is_streaming=False, - ) - request.paged_kv_block_ids = [] - return request - - -def test_context_prompt_lookahead_stops_at_prompt_boundary() -> None: - """An internal chunk hands back the next prompt token; the last one does not. - - The boundary is read off the request's live chunking state rather than a - Python-side prompt length, which a preemption leaves stale. - """ - request = _create_request_with_tokens([10, 11, 12, 13, 14], 1) - - request.context_chunk_size = 2 - assert _get_context_prompt_lookahead_token(request, 2) == 12 - - # A reused prefix advances the live cursor while py_prompt_len remains the - # original prompt length. Internal chunk [2, 4) still looks up token 4. - request.context_current_position = 2 - request.context_chunk_size = 2 - assert _get_context_prompt_lookahead_token(request, 4) == 14 - - # Final chunk [2, 5): nothing follows in the prompt, so the drafter falls - # back to the token the target just sampled. - request.context_chunk_size = 3 - assert (_get_context_prompt_lookahead_token( - request, 5) == INVALID_PROMPT_LOOKAHEAD_TOKEN) - - -def test_zero_fill_row_spans_merges_unsorted_and_adjacent_spans() -> None: - buffer = torch.arange(1, 13, dtype=torch.int32) - _zero_fill_row_spans(buffer, [(9, 11), (2, 4), (4, 5), (7, 8)]) - assert buffer.tolist() == [1, 2, 0, 0, 0, 6, 7, 0, 9, 0, 0, 12] - - -def test_zero_fill_row_spans_without_spans_is_noop() -> None: - buffer = torch.arange(4, dtype=torch.int32) - _zero_fill_row_spans(buffer, []) - assert buffer.tolist() == [0, 1, 2, 3] - - -def _make_request_stub(req_id: int, prompt_len: int = 4) -> SimpleNamespace: - return SimpleNamespace( - py_request_id=req_id, - context_chunk_size=1, - context_remaining_length=1, - context_current_position=prompt_len - 1, - py_prompt_len=prompt_len, - py_beam_width=1, - py_draft_tokens=[], - is_context_only_request=False, - is_generation_only_request=False, - py_disaggregated_params=None, - py_multimodal_data=None, - py_mm_encoder_event=None, - py_mrope_position_delta=None, - py_return_context_logits=False, - py_batch_idx=None, - lora_task_id=None, - is_dummy=False, - max_beam_num_tokens=prompt_len, - state="context", - py_llm_request_type="context_and_generation", - ) - - -def _make_forward_only_engine( - graph_key: KeyType | None, - runner_enabled: bool = True, - private_decoder: bool = False, -) -> tuple[ModelEngine, Mock, Mock, Mock, dict[str, object]]: - engine = object.__new__(ModelEngine) - decoder = object.__new__(DecoderRunner) - engine.model = decoder._model = SimpleNamespace( - extra_attrs={}, - model_config=SimpleNamespace(pretrained_config=SimpleNamespace( - rope_scaling=None))) - engine._warmup_depth = 0 - decoder._config = SimpleNamespace( - compile_only_piecewise_graphs=False, - disable_overlap_scheduler=True, - is_spec_decode=False, - max_beam_width=1, - attention_backend=None, - original_max_draft_len=0, - original_max_total_draft_tokens=0, - spec_dec_max_total_draft_tokens=0, - use_mrope=False, - spec_config=None, - ) - engine.llm_args = SimpleNamespace(mm_encoder_only=False) - engine.mapping = decoder._mapping = SimpleNamespace( - cp_size=1, - enable_lm_head_tp_in_adp=False, - ) - decoder._dist = None - get_runtime_tokens_per_gen_step = Mock(return_value=1) - decoder._get_runtime_tokens_per_gen_step = get_runtime_tokens_per_gen_step - engine.iter_states = decoder._iter_states = {} - decoder.forward_pass_callable = None - decoder.start_beam_histories_callable = None - decoder._stage_in_graph_sampling = None - decoder._moe_load_balancer = None - engine._runner = decoder - decoder._lora = SimpleNamespace(cuda_graph_manager=None) - decoder._trtllm_gen_jit_warmup = False - - semantic_attn_metadata = Mock() - graph_attn_metadata = Mock() - decoder._metadata = SimpleNamespace( - attn_metadata=semantic_attn_metadata, - setup_attention=Mock(return_value=semantic_attn_metadata)) - spec_dec_mode = Mock() - spec_dec_mode.attention_need_spec_dec_mode.return_value = False - spec_dec_mode.is_parallel_draft.return_value = False - spec_metadata = Mock( - spec_dec_mode=spec_dec_mode, - is_spec_dec_tree=False, - is_spec_dec_dynamic_tree=False, - ) - decoder._metadata.spec_metadata = spec_metadata - decoder._metadata.setup_speculative = Mock(return_value=spec_metadata) - prepared_inputs = { - "prepared": True, - "input_ids": torch.zeros(2, dtype=torch.int32), - } - decoder._inputs = SimpleNamespace( - guided_decoder=None, - prepare_inputs=Mock(return_value=PreparedInputs(prepared_inputs)), - record_promoted_iter_states=Mock(), - record_input_copy=Mock()) - outputs = {"logits": object()} - decoder._forward_step = Mock(return_value=outputs) - decoder._execute_logit_post_processors = Mock() - - runner = Mock() - runner.enabled = runner_enabled - runner.pad_batch.side_effect = lambda batch, *_args: nullcontext(batch) - runner.maybe_get_cuda_graph.return_value = ((graph_attn_metadata, None, - graph_key) - if graph_key is not None else - (None, None, None)) - runner.get_graph_pool.return_value = None - runner.needs_capture.return_value = False - runner.is_warmup_only = False - runner.replay.return_value = outputs - decoder._graphs = DecoderCudaGraphs( - cuda_graph_runner=runner, - breakable_cuda_graph_runner=None, - config=decoder._config, - mapping=decoder._mapping, - model=decoder._model, - sparse_attention_config=None, - metrics={}, - dummy_requests=None, - agreement=None, - metadata=decoder._metadata, - get_runtime_tokens_per_gen_step=get_runtime_tokens_per_gen_step, - dynamic_draft_len_mapping=None, - lora=decoder._lora, - forward_warmup=decoder._forward_warmup, - torch_compile_bypass=decoder._bypass_torch_compile, - context_promotion_supported=True, - ) - - resource_manager = Mock() - peft_cache_manager = Mock(data_type=torch.bfloat16) - - def get_resource_manager(resource_type): - if resource_type == ResourceManagerType.PEFT_CACHE_MANAGER: - return peft_cache_manager - if resource_type == ResourceManagerType.KV_CACHE_MANAGER: - return SimpleNamespace(has_private_swa_suffix=Mock( - return_value=private_decoder)) - return object() - - resource_manager.get_resource_manager.side_effect = get_resource_manager - return engine, runner, resource_manager, semantic_attn_metadata, outputs - - -@pytest.mark.cpu_only -@pytest.mark.parametrize("enabled", [False, True]) -@pytest.mark.parametrize( - "ordinary_capture,eligible", - [(False, True), (False, False), (True, True), (True, False)], -) -def test_decoder_forward_routes_compile_bypass(enabled: bool, - ordinary_capture: bool, - eligible: bool) -> None: - """Exercise both the regular forward and ordinary-graph capture callbacks.""" - key = KeyType(batch_size=1, draft_len=0) if ordinary_capture else None - engine, graph_runner, resources, _, outputs = _make_forward_only_engine(key) - eager = torch.nn.Identity() - compiled = torch.nn.Identity() - eager.forward = Mock(return_value=outputs) - compiled.forward = Mock(return_value=outputs) - wrapper = PrefillCompiledModel(eager, compiled) - engine._runner._config.compile_only_piecewise_graphs = enabled - engine._runner._model.model = wrapper - engine._runner._forward_step.side_effect = lambda *args, **kwargs: wrapper() - graph_runner.needs_capture.return_value = ordinary_capture - graph_runner.is_warmup_only = ordinary_capture - graph_runner.capture.side_effect = lambda key, forward, inputs, **kwargs: forward( - inputs) - batch = ScheduledRequests() - batch.generation_requests = [_make_request_stub(1)] - - with patch( - "tensorrt_llm._torch.pyexecutor.model_engine.torch.cuda.Event", - return_value=Mock() - ), patch( - "tensorrt_llm._torch.pyexecutor.engine.runners.decoder.runner.get_per_request_prefill_cuda_graph_flag", - return_value=eligible, - ): - result = engine.forward(batch, resources) - - assert result["logits"] is outputs["logits"] - expected_compiled = eligible and not (enabled and ordinary_capture) - selected = compiled.forward if expected_compiled else eager.forward - other = eager.forward if expected_compiled else compiled.forward - selected.assert_called_once_with() - other.assert_not_called() - assert not wrapper._bypass_active - - -def create_model_engine_and_kvcache( - test_case: unittest.TestCase, - llm_args: TorchLlmArgs | None = None, - execution_stream: torch.cuda.Stream | None = None, - spec_config: DecodingBaseConfig | None = None, - model: DummyModel | None = None, -) -> tuple[ModelEngine, KVCacheManagerV2]: - tokens_per_block = 1 - max_tokens = 258 # Atleast 1 more than the max seq len - num_layers = 1 - batch_size = 13 - - if llm_args is None: - llm_args = TorchLlmArgs(model="dummy", - max_batch_size=batch_size, - max_num_tokens=max_tokens, - cuda_graph_config=CudaGraphConfig( - enable_padding=True, - batch_sizes=[1, 2, 4, 8, 16, 32, 64, 128])) - # The padding tests below rely on these properties of the default - # config: batches of 5 and 13 must round up to 8 and 16 respectively. - test_batches = (5, 13) - for test_batch_size in test_batches: - assert test_batch_size not in llm_args.cuda_graph_config.batch_sizes - - assert (8 in llm_args.cuda_graph_config.batch_sizes - and 16 in llm_args.cuda_graph_config.batch_sizes) - - model_engine = DummyModelEngine(llm_args, - torch.half, - spec_config, - model=model) - - kv_cache_config = KvCacheConfig(max_tokens=max_tokens) - mapping = Mapping(world_size=1, tp_size=1, rank=0) - kv_cache_manager = KVCacheManagerV2( - kv_cache_config, - tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF, - num_layers=num_layers, - num_kv_heads=model_engine.model.config.num_key_value_heads, - head_dim=model_engine.model.config.head_dim, - tokens_per_block=tokens_per_block, - max_seq_len=max_tokens, - max_batch_size=batch_size, - mapping=mapping, - dtype=tensorrt_llm.bindings.DataType.HALF, - execution_stream=execution_stream, - ) - - # Retire GPU caches on the test thread before background GC can collect them. - test_case.addCleanup(kv_cache_manager.shutdown) - return model_engine, kv_cache_manager - - -def _resize_input_buffers(runner: DecoderRunner, max_num_tokens: int) -> None: - """Reallocate the runner's input buffers for ``max_num_tokens`` tokens.""" - buffers = DecoderInputBuffers.for_config( - replace(runner._config, max_num_tokens=max_num_tokens), - runner._use_beam_search) - runner._buffers = runner._inputs._buffers = runner._metadata._buffers = buffers - - -def _set_runner_config(runner: DecoderRunner, **changes) -> None: - """Replace the config shared by the runner and its input components.""" - config = replace(runner._config, **changes) - runner._config = runner._inputs._config = runner._metadata._config = config - - -def _clone_tensors(values: dict) -> dict: - return { - name: value.clone() if isinstance(value, torch.Tensor) else value - for name, value in values.items() - } - - -def _install_graph_runner_stub(runner: DecoderRunner, - key: KeyType | None = None) -> Mock: - """Replace the runner's CUDA graph runner with one that never pads. - - Every lookup hits `stub.key`, or misses while it is None. A hit replays by - recording a copy of the prepared inputs. - """ - stub = Mock() - stub.enabled = True - stub.is_warmup_only = False - stub.key = key - stub.pad_batch.side_effect = lambda batch, *_args: nullcontext(batch) - stub.get_graph_pool.return_value = None - stub.needs_capture.return_value = False - stub.maybe_get_cuda_graph.side_effect = lambda batch, **kwargs: ( - (kwargs["attn_metadata"], kwargs["spec_metadata"], stub.key) - if stub.key is not None else (None, None, None)) - stub.replayed_inputs = [] - - def replay(key, inputs, position_id_offsets=None): - stub.replayed_inputs.append(_clone_tensors(inputs)) - return { - "logits": torch.zeros((inputs["input_ids"].shape[0], 10), - device="cuda") - } - - stub.replay.side_effect = replay - runner._graphs.cuda_graph_runner = stub - return stub - - -@pytest.mark.cpu_only -def test_cleanup_releases_runner_model_before_release_gc() -> None: - - class Model: - pass - - engine = object.__new__(ModelEngine) - model = Model() - model_ref = weakref.ref(model) - engine._cleanup_done = False - engine.model_loader = None - engine._runner = SimpleNamespace(_model=model) - engine._model_caller = ModelCaller(model) - engine._mm_item_scheduler = None - engine._lora_model = model - engine._lora_owner = None - engine.model = model - engine._release_cuda_graphs = Mock() - engine.input_processor = None - del model - - def assert_model_released() -> None: - assert model_ref() is None - - with patch("tensorrt_llm._torch.pyexecutor.model_engine.release_gc", - side_effect=assert_model_released) as release_gc: - engine.cleanup() - - release_gc.assert_called_once_with() - - -class GraphOwnerTeardownTestCase(unittest.TestCase): - - def test_cuda_graph_runner_destructor_does_not_release_window_owner( - self) -> None: - runner = Mock() - - CUDAGraphRunner.__del__(runner) - - runner.clear.assert_called_once_with(release_nccl_window_owners=False) - - def test_engine_graph_ownership_is_split(self) -> None: - clear_blocks_graph = Mock() - model = Mock() - model.modules.return_value = [ - SimpleNamespace(clear_blocks_cuda_graph=clear_blocks_graph) - ] - backend = Mock() - cuda_runner = Mock() - breakable_runner = Mock() - decoder_runner = object.__new__(DecoderRunner) - decoder_runner._graphs = object.__new__(DecoderCudaGraphs) - decoder_runner._graphs.cuda_graph_runner = cuda_runner - decoder_runner._graphs.breakable_cuda_graph_runner = breakable_runner - engine = SimpleNamespace( - model=model, - _torch_compile_backend=backend, - _runner=decoder_runner, - ) - - ModelEngine._release_cuda_graphs(engine, - release_nccl_window_owners=False) - - clear_blocks_graph.assert_not_called() - backend.clear_piecewise_cuda_graphs.assert_called_once_with( - release_nccl_window_owners=False) - for runner in (cuda_runner, breakable_runner): - runner.clear.assert_called_once_with( - release_nccl_window_owners=False) - - ModelEngine._release_model_owned_cuda_graphs( - engine, release_nccl_window_owners=False) - - clear_blocks_graph.assert_called_once_with( - release_nccl_window_owners=False) - - def test_executor_release_clears_model_owned_graphs_first(self) -> None: - for release_model_owned_graphs in (True, False): - with self.subTest( - release_model_owned_graphs=release_model_owned_graphs): - engine = Mock() - - ModelEngine.release_executor_resources( - engine, - release_model_owned_graphs=release_model_owned_graphs) - - expected = [call._release_cuda_graphs()] - if release_model_owned_graphs: - expected.insert(0, call._release_model_owned_cuda_graphs()) - self.assertEqual(engine.method_calls, expected) - - -class SingleTokenContextGraphBatchTestCase(unittest.TestCase): - - def test_generation_only_is_identity(self) -> None: - generation = _make_request_stub(1) - batch = ScheduledRequests() - batch.generation_requests = [generation] - - graph_batch, promoted_ids = _make_single_token_context_graph_batch( - batch) - - self.assertIs(graph_batch, batch) - self.assertEqual(promoted_ids, frozenset()) - - def test_eligible_batch_has_independent_lists_and_stable_order( - self) -> None: - context_0 = _make_request_stub(10, prompt_len=1) - context_1 = _make_request_stub(11, prompt_len=8) - generation = _make_request_stub(12, prompt_len=16) - paused = object() - batch = ScheduledRequests() - batch.context_requests_last_chunk = [context_0, context_1] - batch.generation_requests = [generation] - batch.paused_requests = [paused] - semantic_lists = ( - batch.encoder_requests, - batch.context_requests_chunking, - batch.context_requests_last_chunk, - batch.generation_requests, - batch.paused_requests, - ) - semantic_snapshot = vars(context_1).copy() - - graph_batch, promoted_ids = _make_single_token_context_graph_batch( - batch) - - self.assertIsNot(graph_batch, batch) - self.assertEqual(graph_batch.context_requests, []) - self.assertEqual(graph_batch.generation_requests, - [context_0, context_1, generation]) - self.assertEqual(graph_batch.paused_requests, [paused]) - graph_lists = ( - graph_batch.encoder_requests, - graph_batch.context_requests_chunking, - graph_batch.context_requests_last_chunk, - graph_batch.generation_requests, - graph_batch.paused_requests, - ) - for semantic_list, graph_list in zip(semantic_lists, graph_lists): - self.assertIsNot(semantic_list, graph_list) - self.assertEqual(promoted_ids, frozenset({10, 11})) - self.assertEqual(vars(context_1), semantic_snapshot) - - graph_batch.generation_requests.append(object()) - self.assertEqual(batch.context_requests_last_chunk, - [context_0, context_1]) - self.assertEqual(batch.generation_requests, [generation]) - - def test_structural_fallbacks_return_semantic_batch(self) -> None: - context = _make_request_stub(1) - - encoder_batch = ScheduledRequests() - encoder_batch.encoder_requests = [object()] - encoder_batch.context_requests_last_chunk = [context] - graph_batch, promoted_ids = _make_single_token_context_graph_batch( - encoder_batch) - self.assertIs(graph_batch, encoder_batch) - self.assertFalse(promoted_ids) - - chunking_batch = ScheduledRequests() - chunking_batch.context_requests_chunking = [context] - graph_batch, promoted_ids = _make_single_token_context_graph_batch( - chunking_batch) - self.assertIs(graph_batch, chunking_batch) - self.assertFalse(promoted_ids) - - def test_context_shape_and_mode_fallback_matrix(self) -> None: - cases = ( - ("multi_token_chunk", "context_chunk_size", 2), - ("more_context_remaining", "context_remaining_length", 2), - ("cursor_prompt_mismatch", "py_prompt_len", 5), - ("beam", "py_beam_width", 2), - ("draft", "py_draft_tokens", [9]), - ("context_only", "is_context_only_request", True), - ("disaggregated", "py_disaggregated_params", object()), - ("multimodal", "py_multimodal_data", {}), - ("multimodal_event", "py_mm_encoder_event", object()), - ) - for name, attribute, value in cases: - with self.subTest(name=name): - context = _make_request_stub(1) - setattr(context, attribute, value) - batch = ScheduledRequests() - batch.context_requests_last_chunk = [context] - - graph_batch, promoted_ids = \ - _make_single_token_context_graph_batch(batch) - - self.assertIs(graph_batch, batch) - self.assertFalse(promoted_ids) - - def test_context_logits_use_final_token_graph_candidate(self) -> None: - context = _make_request_stub(1) - context.py_return_context_logits = True - batch = ScheduledRequests() - batch.context_requests_last_chunk = [context] - - graph_batch, promoted_ids = _make_single_token_context_graph_batch( - batch) - - self.assertIsNot(graph_batch, batch) - self.assertEqual(graph_batch.generation_requests, [context]) - self.assertEqual(promoted_ids, frozenset({context.py_request_id})) - - def test_generation_only_request_in_context_list_falls_back(self) -> None: - context = _make_request_stub(1) - context.is_generation_only_request = True - batch = ScheduledRequests() - batch.context_requests_last_chunk = [context] - graph_batch, promoted_ids = _make_single_token_context_graph_batch( - batch) - self.assertIs(graph_batch, batch) - self.assertFalse(promoted_ids) - - def test_generation_shape_fallback_matrix(self) -> None: - cases = ( - ("beam", "py_beam_width", 2), - ("draft", "py_draft_tokens", [9]), - ("disaggregated", "py_disaggregated_params", object()), - ) - for name, attribute, value in cases: - with self.subTest(name=name): - context = _make_request_stub(1) - generation = _make_request_stub(2) - setattr(generation, attribute, value) - batch = ScheduledRequests() - batch.context_requests_last_chunk = [context] - batch.generation_requests = [generation] - - graph_batch, promoted_ids = \ - _make_single_token_context_graph_batch(batch) - - self.assertIs(graph_batch, batch) - self.assertFalse(promoted_ids) - - def test_mixed_one_and_two_token_contexts_fall_back_together(self) -> None: - one_token = _make_request_stub(1) - two_tokens = _make_request_stub(2) - two_tokens.context_current_position -= 1 - two_tokens.context_remaining_length = 2 - two_tokens.context_chunk_size = 2 - batch = ScheduledRequests() - batch.context_requests_last_chunk = [one_token, two_tokens] - - graph_batch, promoted_ids = _make_single_token_context_graph_batch( - batch) - - self.assertIs(graph_batch, batch) - self.assertFalse(promoted_ids) - self.assertEqual(batch.context_requests_last_chunk, - [one_token, two_tokens]) - - def test_mrope_delta_is_supported_by_decode_provider(self) -> None: - context = _make_request_stub(1) - context.py_mrope_position_delta = object() - batch = ScheduledRequests() - batch.context_requests_last_chunk = [context] - - graph_batch, promoted_ids = _make_single_token_context_graph_batch( - batch) - - self.assertIsNot(graph_batch, batch) - self.assertEqual(graph_batch.generation_requests, [context]) - self.assertEqual(promoted_ids, frozenset({context.py_request_id})) - - def test_multimodal_context_requires_compatible_decode_token(self) -> None: - context = _make_request_stub(1) - context.py_multimodal_data = {} - batch = ScheduledRequests() - batch.context_requests_last_chunk = [context] - - graph_batch, promoted_ids = _make_single_token_context_graph_batch( - batch) - self.assertIs(graph_batch, batch) - self.assertFalse(promoted_ids) - - incompatible = Mock(return_value=False) - graph_batch, promoted_ids = _make_single_token_context_graph_batch( - batch, incompatible) - self.assertIs(graph_batch, batch) - self.assertFalse(promoted_ids) - incompatible.assert_called_once_with(context) - - compatible = Mock(return_value=True) - graph_batch, promoted_ids = _make_single_token_context_graph_batch( - batch, compatible) - self.assertIsNot(graph_batch, batch) - self.assertEqual(graph_batch.generation_requests, [context]) - self.assertEqual(promoted_ids, frozenset({context.py_request_id})) - compatible.assert_called_once_with(context) - - def test_multimodal_pending_event_is_rechecked(self) -> None: - context = _make_request_stub(1) - context.py_multimodal_data = {} - context.py_mm_encoder_event = object() - batch = ScheduledRequests() - batch.context_requests_last_chunk = [context] - compatible = Mock(return_value=True) - - graph_batch, promoted_ids = _make_single_token_context_graph_batch( - batch, compatible) - self.assertIs(graph_batch, batch) - self.assertFalse(promoted_ids) - compatible.assert_not_called() - - context.py_mm_encoder_event = None - graph_batch, promoted_ids = _make_single_token_context_graph_batch( - batch, compatible) - self.assertIsNot(graph_batch, batch) - self.assertEqual(promoted_ids, frozenset({context.py_request_id})) - compatible.assert_called_once_with(context) - - def test_multimodal_decode_compatibility_uses_final_prompt_token( - self) -> None: - graphs = object.__new__(DecoderCudaGraphs) - graphs.model = SimpleNamespace( - config=SimpleNamespace(vocab_size=100), - mm_token_ids=torch.tensor([99], dtype=torch.int32), - ) - graphs._config = SimpleNamespace(use_mrope=False) - request = _create_request_with_tokens([11, 99, 22], 1) - - request.context_current_position = 2 - self.assertTrue( - graphs._is_final_multimodal_context_decode_compatible(request)) - - request.context_current_position = 1 - self.assertFalse( - graphs._is_final_multimodal_context_decode_compatible(request)) - - graphs.model.mm_token_ids = None - request = _create_request_with_tokens([11, 100], 2) - request.context_current_position = 1 - self.assertFalse( - graphs._is_final_multimodal_context_decode_compatible(request)) - - graphs.model.mm_token_ids = torch.tensor([99], dtype=torch.int32) - graphs.model.model_config = SimpleNamespace( - pretrained_config=SimpleNamespace(rope_scaling={"type": "mrope"})) - graphs._config.use_mrope = True - request = _create_request_with_tokens([11, 22], 3) - request.context_current_position = 1 - request.py_multimodal_data = {"mrope_config": {}} - self.assertTrue( - graphs._is_final_multimodal_context_decode_compatible(request)) - - request.py_multimodal_data["multimodal_embedding"] = object() - self.assertFalse( - graphs._is_final_multimodal_context_decode_compatible(request)) - - request.py_multimodal_data["mrope_config"][ - "mrope_position_deltas"] = object() - self.assertTrue( - graphs._is_final_multimodal_context_decode_compatible(request)) - - def test_sparse_sequence_mode_uses_promoted_context_cursor(self) -> None: - sparse_config = Mock(spec=SeqLenAwareSparseAttentionConfig) - sparse_config.needs_separate_short_long_cuda_graphs.return_value = True - sparse_config.seq_len_threshold = 16 - runner = object.__new__(CUDAGraphRunner) - runner.sparse_config = sparse_config - runner.spec_config = None - runner.graphs = {} - runner.graph_outputs = {} - runner.graph_metadata = {} - runner.padding_dummy_requests = {} - runner.memory_pool = None - - request = _make_request_stub(7, prompt_len=8) - request.py_batch_idx = 0 - request.max_beam_num_tokens = 64 - batch = ScheduledRequests() - batch.generation_requests = [request] - overlap_state = SimpleNamespace(new_tokens=object()) - - self.assertTrue( - runner._get_seq_len_mode(batch, overlap_state, - frozenset({request.py_request_id}))) - self.assertFalse( - runner._get_seq_len_mode(batch, overlap_state, frozenset())) - - def test_graph_key_forwards_promoted_context_ids(self) -> None: - runner = Mock() - runner.config = SimpleNamespace() - runner._resolve_sample_type.return_value = SampleType.FULL - runner._get_seq_len_mode.return_value = True - runner._ragged_verify_bucket.return_value = None - request = _make_request_stub(7) - batch = ScheduledRequests() - batch.generation_requests = [request] - promoted_ids = frozenset({request.py_request_id}) - - key = CUDAGraphRunner.get_graph_key( - runner, - batch, - new_tensors_device=None, - promoted_context_request_ids=promoted_ids, - ) - - runner._get_seq_len_mode.assert_called_once_with( - batch, None, promoted_ids) - self.assertEqual( - key, - KeyType(batch_size=1, - draft_len=0, - short_seq_len_mode=True, - num_encoder_tokens=0)) - - def test_graph_key_aggregates_encoder_tokens(self) -> None: - runner = Mock() - runner.config = SimpleNamespace() - runner._resolve_sample_type.return_value = SampleType.FULL - runner.max_beam_width = 1 - runner._get_seq_len_mode.return_value = False - runner._ragged_verify_bucket.return_value = None - context = _make_request_stub(1) - context.encoder_output_len = 7 - context.py_skip_cross_kv_projection = False - skipped_context = _make_request_stub(2) - skipped_context.encoder_output_len = 11 - skipped_context.py_skip_cross_kv_projection = True - batch = ScheduledRequests() - batch.context_requests_last_chunk = [context, skipped_context] - batch.generation_requests = [_make_request_stub(3)] - - key = CUDAGraphRunner.get_graph_key(runner, batch) - - assert key is not None - self.assertEqual(key.num_contexts, 2) - self.assertEqual(key.context_query_len, 1) - self.assertEqual(key.num_encoder_tokens, 7) - self.assertEqual(CUDAGraphRunner._get_num_tokens_for_key(runner, key), - 3) - - def test_graph_key_rejects_nonuniform_context_query_lengths(self) -> None: - runner = Mock() - runner.config = SimpleNamespace() - runner._resolve_sample_type.return_value = SampleType.FULL - runner._get_seq_len_mode.return_value = False - runner._ragged_verify_bucket.return_value = None - first_context = _make_request_stub(1) - first_context.encoder_output_len = 7 - first_context.py_skip_cross_kv_projection = False - second_context = _make_request_stub(2) - second_context.context_chunk_size = 2 - second_context.encoder_output_len = 11 - second_context.py_skip_cross_kv_projection = False - batch = ScheduledRequests() - batch.context_requests_last_chunk = [first_context, second_context] - batch.generation_requests = [_make_request_stub(3)] - - key = CUDAGraphRunner.get_graph_key(runner, batch) - - self.assertIsNone(key) - - def test_graph_key_rounds_encoder_tokens_up_to_captured_extent( - self) -> None: - key = KeyType(batch_size=2, - draft_len=0, - num_contexts=1, - context_query_len=1, - num_encoder_tokens=7) - smaller_key = key._replace(num_encoder_tokens=6) - compatible_key = key._replace(num_encoder_tokens=8) - larger_key = key._replace(num_encoder_tokens=16) - runner = SimpleNamespace( - padding_enabled=True, - _capture_allowed=False, - graph_metadata={}, - graph_outputs={ - smaller_key: object(), - compatible_key: object(), - larger_key: object(), - }, - ) - - actual_key = CUDAGraphRunner._get_compatible_mixed_encoder_decoder_key( - runner, key) - - self.assertEqual(actual_key, compatible_key) - - def test_graph_key_includes_peft_cache_dtype(self) -> None: - runner = Mock() - runner.config = SimpleNamespace() - runner._resolve_sample_type.return_value = SampleType.FULL - runner._get_seq_len_mode.return_value = False - runner._ragged_verify_bucket.return_value = None - request = _make_request_stub(7) - batch = ScheduledRequests() - batch.generation_requests = [request] - - model_dtype_key = CUDAGraphRunner.get_graph_key( - runner, batch, peft_cache_data_type=torch.bfloat16) - fp8_key = CUDAGraphRunner.get_graph_key( - runner, batch, peft_cache_data_type=torch.float8_e4m3fn) - - self.assertNotEqual(model_dtype_key, fp8_key) - self.assertEqual( - model_dtype_key._replace(peft_cache_data_type=None), - fp8_key._replace(peft_cache_data_type=None), - ) - - def test_graph_dtype_change_falls_back_to_eager(self) -> None: - runner = Mock() - runner.enabled = True - runner._agree_graph_shape = CUDAGraphRunner._agree_graph_shape.__get__( - runner) - runner._dspark_confidence_enabled = False - runner.config = SimpleNamespace( - enable_attention_dp=False, - use_mrope=False, - ) - model_dtype_key = KeyType(batch_size=1, - draft_len=0, - peft_cache_data_type=torch.bfloat16) - fp8_key = KeyType(batch_size=1, - draft_len=0, - peft_cache_data_type=torch.float8_e4m3fn) - runner.get_graph_key.return_value = fp8_key - runner.graph_metadata = {model_dtype_key: object()} - runner._capture_allowed = False - runner._is_mixed_encoder_decoder_batch.return_value = False - runner._can_run_cuda_graph_batch.return_value = True - - request = _make_request_stub(7) - batch = ScheduledRequests() - batch.generation_requests = [request] - - with patch( - "tensorrt_llm._torch.pyexecutor.cuda_graph_runner.ExpertStatistic.should_record", - return_value=False): - result = CUDAGraphRunner.maybe_get_cuda_graph( - runner, - batch, - enable_spec_decode=False, - attn_metadata=object(), - peft_cache_data_type=torch.float8_e4m3fn, - ) - - self.assertEqual(result, (None, None, None)) - - def test_graph_key_includes_lora_variant(self) -> None: - runner = Mock() - runner.config = SimpleNamespace() - runner._resolve_sample_type.return_value = SampleType.FULL - runner._get_seq_len_mode.return_value = False - runner._ragged_verify_bucket.return_value = None - request = _make_request_stub(7) - batch = ScheduledRequests() - batch.generation_requests = [request] - - key = CUDAGraphRunner.get_graph_key( - runner, - batch, - use_lora_graph=True, - ) - - self.assertEqual( - key, - KeyType(batch_size=1, draft_len=0, use_lora_graph=True), - ) - - def test_lora_graph_variant_selection(self) -> None: - graphs = object.__new__(DecoderCudaGraphs) - graphs._lora = SimpleNamespace(cuda_graph_manager=object()) - graphs._force_lora_graph_for_capture = None - graphs._config = SimpleNamespace(cuda_graph_specialize_lora=True) - request = _make_request_stub(7) - batch = ScheduledRequests() - batch.generation_requests = [request] - - self.assertFalse(graphs.use_lora_cuda_graph(batch)) - - request.lora_task_id = 42 - self.assertTrue(graphs.use_lora_cuda_graph(batch)) - - request.lora_task_id = None - graphs._config.cuda_graph_specialize_lora = False - self.assertTrue(graphs.use_lora_cuda_graph(batch)) - - graphs._force_lora_graph_for_capture = False - self.assertFalse(graphs.use_lora_cuda_graph(batch)) - - def test_graph_lookup_forwards_promoted_context_ids(self) -> None: - runner = Mock() - runner.enabled = True - runner._agree_graph_shape = CUDAGraphRunner._agree_graph_shape.__get__( - runner) - runner._dspark_confidence_enabled = False - runner.config = SimpleNamespace( - enable_attention_dp=False, - use_mrope=False, - ) - key = KeyType(batch_size=1, - draft_len=0, - short_seq_len_mode=True, - num_encoder_tokens=0) - graph_attn_metadata = object() - graph_spec_metadata = object() - runner.get_graph_key.return_value = key - runner.graphs = {key: object()} - runner.graph_metadata = { - key: { - "attn_metadata": graph_attn_metadata, - "spec_metadata": graph_spec_metadata, - } - } - runner._is_mixed_encoder_decoder_batch.return_value = False - request = _make_request_stub(7) - batch = ScheduledRequests() - batch.generation_requests = [request] - promoted_ids = frozenset({request.py_request_id}) - - with patch( - "tensorrt_llm._torch.pyexecutor.cuda_graph_runner.ExpertStatistic.should_record", - return_value=False): - result = CUDAGraphRunner.maybe_get_cuda_graph( - runner, - batch, - enable_spec_decode=False, - attn_metadata=object(), - promoted_context_request_ids=promoted_ids, - ) - - runner.get_graph_key.assert_called_once_with(batch, None, None, None, - promoted_ids, None, False, - None) - self.assertEqual(result, - (graph_attn_metadata, graph_spec_metadata, key)) - - def test_ced_final_context_cannot_promote_to_decode_graph(self) -> None: - key = KeyType(batch_size=1, draft_len=0) - engine, runner, resource_manager, metadata, outputs = \ - _make_forward_only_engine(key, private_decoder=True) - engine.model.model = SimpleNamespace(ced_kv_precompute=True, - decoder_replay_split=20) - batch = ScheduledRequests() - batch.context_requests_last_chunk = [_make_request_stub(1)] - with patch( - "tensorrt_llm._torch.pyexecutor.model_engine.torch.cuda.Event" - ), patch( - "tensorrt_llm._torch.pyexecutor.engine.runners.decoder.graphs._make_single_token_context_graph_batch" - ) as promote: - actual_outputs = engine.forward(batch, resource_manager) - self.assertEqual(actual_outputs, { - **outputs, "runtime_draft_len": 0 - }) - self.assertIs(actual_outputs["logits"], outputs["logits"]) - promote.assert_not_called() - runner.replay.assert_not_called() - engine._runner._forward_step.assert_called_once() - self.assertIs(engine._runner._inputs.prepare_inputs.call_args.args[2], - metadata) - - def test_zero_runtime_draft_speculation_commits_graph_candidate( - self) -> None: - key = KeyType(batch_size=2, draft_len=0) - engine, runner, resource_manager, semantic_attn_metadata, outputs = \ - _make_forward_only_engine(key) - engine._runner._config.spec_config = SimpleNamespace( - spec_dec_mode=_TWO_MODEL_SPEC_MODE, - is_linear_tree=True, - get_runtime_tokens_per_gen_step=Mock(return_value=1)) - graph_attn_metadata = runner.maybe_get_cuda_graph.return_value[0] - runner.maybe_get_cuda_graph.return_value = ( - graph_attn_metadata, - engine.spec_metadata, - key, - ) - context = _make_request_stub(1) - generation = _make_request_stub(2) - batch = ScheduledRequests() - batch.context_requests_last_chunk = [context] - batch.generation_requests = [generation] - - with patch( - "tensorrt_llm._torch.pyexecutor.model_engine.torch.cuda.Event", - return_value=Mock()): - actual_outputs = engine.forward(batch, - resource_manager, - enable_spec_decode=True) - - self.assertEqual(actual_outputs, {**outputs, "runtime_draft_len": 0}) - self.assertIsNot(actual_outputs, outputs) - graph_batch = runner.maybe_get_cuda_graph.call_args.args[0] - self.assertEqual(graph_batch.generation_requests, [context, generation]) - self.assertTrue( - runner.maybe_get_cuda_graph.call_args.kwargs["enable_spec_decode"]) - engine.spec_metadata.update_is_all_greedy_sample.assert_called_once_with( - graph_batch.all_requests()) - prepare_args = engine._runner._inputs.prepare_inputs.call_args.args - self.assertIs(prepare_args[0], graph_batch) - self.assertIs(prepare_args[3], engine.spec_metadata) - self.assertEqual(prepare_args[-1], frozenset({context.py_request_id})) - semantic_attn_metadata.update_spec_dec_param.assert_called_once() - self.assertEqual( - semantic_attn_metadata.update_spec_dec_param.call_args. - kwargs["num_contexts"], 1) - prepared_inputs = engine._runner._inputs.prepare_inputs.return_value.kwargs - runner.replay.assert_called_once_with(key, - prepared_inputs, - position_id_offsets=None) - - def test_zero_runtime_draft_speculation_graph_miss_is_semantic_eager( - self) -> None: - engine, runner, resource_manager, semantic_attn_metadata, outputs = \ - _make_forward_only_engine(None) - engine._runner._config.spec_config = SimpleNamespace( - spec_dec_mode=_TWO_MODEL_SPEC_MODE, - is_linear_tree=True, - get_runtime_tokens_per_gen_step=Mock(return_value=1)) - context = _make_request_stub(1) - batch = ScheduledRequests() - batch.context_requests_last_chunk = [context] - - with patch( - "tensorrt_llm._torch.pyexecutor.model_engine.torch.cuda.Event", - return_value=Mock()): - actual_outputs = engine.forward(batch, - resource_manager, - enable_spec_decode=True) - - self.assertEqual(actual_outputs, {**outputs, "runtime_draft_len": 0}) - self.assertIsNot(actual_outputs, outputs) - graph_batch = runner.maybe_get_cuda_graph.call_args.args[0] - self.assertEqual(graph_batch.generation_requests, [context]) - prepare_args = engine._runner._inputs.prepare_inputs.call_args.args - self.assertIs(prepare_args[0], batch) - self.assertIs(prepare_args[2], semantic_attn_metadata) - self.assertIs(prepare_args[3], engine.spec_metadata) - self.assertEqual(prepare_args[-1], frozenset()) - engine._runner._forward_step.assert_called_once() - runner.replay.assert_not_called() - - def test_zero_runtime_non_linear_tree_speculation_uses_semantic_eager_batch( - self) -> None: - engine, runner, resource_manager, semantic_attn_metadata, outputs = \ - _make_forward_only_engine(None) - engine._runner._config.spec_config = SimpleNamespace( - spec_dec_mode=_TWO_MODEL_SPEC_MODE, - is_linear_tree=False, - tokens_per_gen_step=8, - get_runtime_tokens_per_gen_step=Mock(return_value=1)) - context = _make_request_stub(1) - generation = _make_request_stub(2) - batch = ScheduledRequests() - batch.context_requests_last_chunk = [context] - batch.generation_requests = [generation] - - with patch( - "tensorrt_llm._torch.pyexecutor.model_engine.torch.cuda.Event", - return_value=Mock()): - actual_outputs = engine.forward(batch, - resource_manager, - enable_spec_decode=True) - - self.assertEqual(actual_outputs, {**outputs, "runtime_draft_len": 0}) - self.assertIsNot(actual_outputs, outputs) - lookup = runner.maybe_get_cuda_graph.call_args - self.assertIs(lookup.args[0], batch) - self.assertEqual(lookup.kwargs["promoted_context_request_ids"], - frozenset()) - prepare_args = engine._runner._inputs.prepare_inputs.call_args.args - self.assertIs(prepare_args[0], batch) - self.assertIs(prepare_args[2], semantic_attn_metadata) - self.assertIs(prepare_args[3], engine.spec_metadata) - self.assertEqual(prepare_args[-1], frozenset()) - engine._runner._forward_step.assert_called_once() - runner.replay.assert_not_called() - - def test_multimodal_graph_miss_preserves_semantic_payload(self) -> None: - engine, runner, resource_manager, _, _ = _make_forward_only_engine(None) - engine.model.config = SimpleNamespace(vocab_size=100) - engine.model.mm_token_ids = torch.tensor([99], dtype=torch.int32) - context = _make_request_stub(1, prompt_len=3) - context.get_tokens = Mock(return_value=[99, 11, 22]) - multimodal_data = { - "multimodal_embedding": object(), - "mrope_config": { - "mrope_position_deltas": object() - }, - } - context.py_multimodal_data = multimodal_data - batch = ScheduledRequests() - batch.context_requests_last_chunk = [context] - - with patch( - "tensorrt_llm._torch.pyexecutor.model_engine.torch.cuda.Event", - return_value=Mock()): - engine.forward(batch, resource_manager) - - graph_batch = runner.maybe_get_cuda_graph.call_args.args[0] - self.assertIsNot(graph_batch, batch) - self.assertEqual(graph_batch.generation_requests, [context]) - self.assertIs(engine._runner._inputs.prepare_inputs.call_args.args[0], - batch) - self.assertIs(context.py_multimodal_data, multimodal_data) - self.assertIn("multimodal_embedding", multimodal_data) - - def test_breakable_graph_falls_back_for_context_logits(self) -> None: - for return_context_logits in (True, False): - with self.subTest(return_context_logits=return_context_logits): - engine, _, resource_manager, _, outputs = \ - _make_forward_only_engine(None) - breakable_runner = Mock() - breakable_runner.is_capturing = False - breakable_runner.is_warming_up = False - breakable_runner.has_graph.return_value = True - breakable_runner.execute.return_value = outputs - engine._runner._graphs.breakable_cuda_graph_runner = breakable_runner - prepared = engine._runner._inputs.prepare_inputs.return_value - - def prepare_prefill_graph_batch(*args, **kwargs): - set_per_request_prefill_cuda_graph_flag(True) - return prepared - - engine._runner._inputs.prepare_inputs.side_effect = \ - prepare_prefill_graph_batch - self.addCleanup(set_per_request_prefill_cuda_graph_flag, False) - - batch = ScheduledRequests() - batch.context_requests_last_chunk = [_make_request_stub(1)] - batch.context_requests_last_chunk[ - 0].py_return_context_logits = return_context_logits - - with patch( - "tensorrt_llm._torch.pyexecutor.model_engine.torch.cuda.Event", - return_value=Mock()): - actual_outputs = engine.forward( - batch, - resource_manager, - ) - - self.assertEqual(actual_outputs, { - **outputs, "runtime_draft_len": 0 - }) - self.assertIsNot(actual_outputs, outputs) - if return_context_logits: - breakable_runner.execute.assert_not_called() - engine._runner._forward_step.assert_called_once() - else: - breakable_runner.execute.assert_called_once() - engine._runner._forward_step.assert_not_called() - - def test_generation_only_forward_does_not_call_new_selector(self) -> None: - key = KeyType(batch_size=1, draft_len=0) - engine, runner, resource_manager, _, _ = _make_forward_only_engine(key) - generation = _make_request_stub(2) - batch = ScheduledRequests() - batch.generation_requests = [generation] - - with patch( - "tensorrt_llm._torch.pyexecutor.model_engine.torch.cuda.Event", - return_value=Mock()): - engine.forward(batch, resource_manager) - - lookup = runner.maybe_get_cuda_graph.call_args - self.assertIs(lookup.args[0], batch) - self.assertEqual(lookup.kwargs["promoted_context_request_ids"], - frozenset()) - self.assertFalse(lookup.kwargs["use_lora_graph"]) - self.assertIs(engine._runner._inputs.prepare_inputs.call_args.args[0], - batch) - self.assertEqual( - engine._runner._inputs.prepare_inputs.call_args.args[-1], - frozenset()) - - def test_generation_lora_request_selects_lora_graph(self) -> None: - key = KeyType(batch_size=1, draft_len=0, use_lora_graph=True) - engine, runner, resource_manager, _, _ = _make_forward_only_engine(key) - engine._runner._lora = engine._runner._graphs._lora = SimpleNamespace( - cuda_graph_manager=object()) - engine._runner._config.cuda_graph_specialize_lora = True - generation = _make_request_stub(2) - generation.lora_task_id = 42 - batch = ScheduledRequests() - batch.generation_requests = [generation] - - with patch( - "tensorrt_llm._torch.pyexecutor.model_engine.torch.cuda.Event", - return_value=Mock()): - engine.forward(batch, resource_manager) - - self.assertTrue( - runner.maybe_get_cuda_graph.call_args.kwargs["use_lora_graph"]) - self.assertEqual( - runner.maybe_get_cuda_graph.call_args. - kwargs["peft_cache_data_type"], torch.bfloat16) - self.assertTrue(engine._runner._inputs.prepare_inputs.call_args. - kwargs["use_lora_graph"]) - - def test_global_incompatibilities_bypass_candidate_selection(self) -> None: - cases = ( - "graphs_disabled", - "speculative_nonzero_runtime_draft", - "beam", - "extra_inputs", - "ple_recurrent_state", - "nested_ple_recurrent_state", - "context_parallel", - "context_only_cache", - ) - for case in cases: - with self.subTest(case=case): - engine, runner, resource_manager, _, _ = \ - _make_forward_only_engine(None) - spec_state = {} - if case == "graphs_disabled": - runner.enabled = False - elif case == "speculative_nonzero_runtime_draft": - spec_state = dict(enable_spec_decode=True, - runtime_draft_len=1) - engine._runner._config.spec_config = SimpleNamespace( - spec_dec_mode=_TWO_MODEL_SPEC_MODE, - is_linear_tree=True, - get_runtime_tokens_per_gen_step=Mock(return_value=1)) - elif case == "beam": - engine._runner._config.max_beam_width = 2 - elif case == "extra_inputs": - engine._runner._graphs._context_promotion_supported = False - elif case == "ple_recurrent_state": - engine.model.has_ple = True - elif case == "nested_ple_recurrent_state": - engine.model.model = SimpleNamespace(llm=SimpleNamespace( - model=SimpleNamespace(has_ple=True))) - elif case == "context_parallel": - engine.mapping.cp_size = 2 - elif case == "context_only_cache": - resource_manager.get_resource_manager.side_effect = ( - lambda _: SimpleNamespace(context_swa_layer_limit=20)) - - batch = ScheduledRequests() - batch.context_requests_last_chunk = [_make_request_stub(1)] - with patch( - "tensorrt_llm._torch.pyexecutor.model_engine.torch.cuda.Event", - return_value=Mock()): - engine.forward(batch, resource_manager, **spec_state) - - lookup = runner.maybe_get_cuda_graph.call_args - self.assertIs(lookup.args[0], batch) - self.assertEqual(lookup.kwargs["promoted_context_request_ids"], - frozenset()) - self.assertIs( - engine._runner._inputs.prepare_inputs.call_args.args[0], - batch) - - def test_forward_threads_call_state_and_returns_prepared_length( - self) -> None: - engine, runner, resource_manager, _, outputs = \ - _make_forward_only_engine(None) - engine._runner._config.spec_config = SimpleNamespace( - spec_dec_mode=_TWO_MODEL_SPEC_MODE, - is_linear_tree=False, - tokens_per_gen_step=8, - get_runtime_tokens_per_gen_step=Mock(return_value=1)) - # A non-linear tree widens the length to its total draft tokens during - # preparation, after graph selection. - prepared = engine._runner._inputs.prepare_inputs.return_value - engine._runner._inputs.prepare_inputs.return_value = replace( - prepared, runtime_draft_len=7) - batch = ScheduledRequests() - batch.generation_requests = [_make_request_stub(2)] - - with patch( - "tensorrt_llm._torch.pyexecutor.model_engine.torch.cuda.Event", - return_value=Mock()): - actual_outputs = engine.forward(batch, - resource_manager, - enable_spec_decode=True, - runtime_draft_len=2) - - self.assertEqual(actual_outputs, {**outputs, "runtime_draft_len": 7}) - self.assertEqual(runner.pad_batch.call_args.args[2], 2) - self.assertTrue( - runner.maybe_get_cuda_graph.call_args.kwargs["enable_spec_decode"]) - prepare_kwargs = engine._runner._inputs.prepare_inputs.call_args.kwargs - self.assertEqual( - (prepare_kwargs["enable_spec_decode"], - prepare_kwargs["runtime_draft_len"], prepare_kwargs["is_dummy"]), - (True, 2, False)) - step_kwargs = engine._runner._forward_step.call_args.kwargs - self.assertEqual( - (step_kwargs["enable_spec_decode"], - step_kwargs["runtime_draft_len"], step_kwargs["is_dummy"]), - (True, 7, False)) - - -class DecoderForwardTestCase(unittest.TestCase): - """Drive the decoder runner through `engine.forward`. - - The runner, KV cache manager and requests are real; only the CUDA graph - runner is stubbed. Assertions read what the model or graph replay received. - """ - - def _make_engine( - self, - model: DummyModel | None = None, - ) -> tuple[ModelEngine, KVCacheManagerV2, ResourceManager, Mock]: - engine, kv_cache_manager = create_model_engine_and_kvcache(self, - model=model) - resource_manager = ResourceManager( - {ResourceManagerType.KV_CACHE_MANAGER: kv_cache_manager}) - graph_runner = _install_graph_runner_stub(engine._runner) - return engine, kv_cache_manager, resource_manager, graph_runner - - @staticmethod - def _forward(engine: ModelEngine, - kv_cache_manager: KVCacheManagerV2, - resource_manager: ResourceManager, - batch: ScheduledRequests, - new_tensors_device=None): - _prepare_kv_cache_resources(kv_cache_manager, batch) - return engine.forward(batch, resource_manager, new_tensors_device) - - @staticmethod - def _request(tokens: list[int], req_id: int, seq_slot: int) -> LlmRequest: - request = _create_request_with_tokens(tokens, req_id) - request.py_seq_slot = seq_slot - return request - - def _prefill(self, engine, kv_cache_manager, resource_manager, *requests: - LlmRequest) -> None: - batch = ScheduledRequests() - batch.context_requests_last_chunk = list(requests) - self._forward(engine, kv_cache_manager, resource_manager, batch) - - def _prefill_all_but_last_token(self, engine, kv_cache_manager, - resource_manager, - request: LlmRequest) -> None: - """Leave `request` with a final one-token context chunk.""" - batch = ScheduledRequests() - batch.context_requests_chunking = [request] - _prepare_kv_cache_resources(kv_cache_manager, batch) - # Like the scheduler, narrow the chunk after V2 spans it to the prompt end. - request.context_chunk_size = request.py_prompt_len - 1 - engine.forward(batch, resource_manager) - request.move_to_next_context_chunk() - request.context_chunk_size = 1 - - def _final_context_and_generation(self, engine, kv_cache_manager, - resource_manager): - context = self._request([11, 22, 33, 44], 1, seq_slot=0) - generation = self._request([50, 51, 52, 53, 54], 2, seq_slot=1) - self._prefill_all_but_last_token(engine, kv_cache_manager, - resource_manager, context) - self._prefill(engine, kv_cache_manager, resource_manager, generation) - generation.add_new_token(60, 0) - batch = ScheduledRequests() - batch.context_requests_last_chunk = [context] - batch.generation_requests = [generation] - return context, generation, batch - - def test_forward_commits_candidate_only_on_graph_hit(self) -> None: - engine, kv_cache_manager, resource_manager, graph_runner = \ - self._make_engine() - context, generation, batch = self._final_context_and_generation( - engine, kv_cache_manager, resource_manager) - num_model_calls = len(engine.model.recorded_kwargs) - graph_runner.key = KeyType(batch_size=2, draft_len=0) - - outputs = self._forward(engine, kv_cache_manager, resource_manager, - batch) - - lookup = graph_runner.maybe_get_cuda_graph.call_args - self.assertIsNot(lookup.args[0], batch) - self.assertEqual(lookup.args[0].generation_requests, - [context, generation]) - self.assertEqual(lookup.kwargs["promoted_context_request_ids"], - frozenset({context.py_request_id})) - graph_runner.replay.assert_called_once() - self.assertEqual(graph_runner.replay.call_args.args[0], - graph_runner.key) - replayed = graph_runner.replayed_inputs[0] - self.assertEqual(replayed["input_ids"].tolist(), [44, 60]) - self.assertEqual(replayed["position_ids"].flatten().tolist(), [3, 5]) - self.assertEqual(len(engine.model.recorded_kwargs), num_model_calls) - self.assertEqual(set(outputs), {"logits", "runtime_draft_len"}) - self.assertEqual(engine.attn_metadata.num_contexts, 0) - self.assertEqual(engine.iter_states['num_ctx_requests'], 1) - self.assertEqual(engine.iter_states['num_ctx_tokens'], 1) - self.assertEqual(engine.iter_states['num_generation_tokens'], 1) - - def test_forward_graph_miss_uses_semantic_eager_batch(self) -> None: - engine, kv_cache_manager, resource_manager, graph_runner = \ - self._make_engine() - context, generation, batch = self._final_context_and_generation( - engine, kv_cache_manager, resource_manager) - - self._forward(engine, kv_cache_manager, resource_manager, batch) - - lookup = graph_runner.maybe_get_cuda_graph.call_args - self.assertEqual(lookup.args[0].generation_requests, - [context, generation]) - graph_runner.replay.assert_not_called() - model_kwargs = engine.model.recorded_kwargs[-1] - self.assertEqual(model_kwargs["input_ids"].tolist(), [44, 60]) - self.assertEqual(model_kwargs["position_ids"].flatten().tolist(), - [3, 5]) - self.assertEqual(engine.attn_metadata.num_contexts, 1) - - def test_forward_allows_guided_context_logits_on_graph_hit(self) -> None: - engine, kv_cache_manager, resource_manager, graph_runner = \ - self._make_engine() - engine._runner.guided_decoder = Mock() - context = self._request([11, 22, 33, 44], 1, seq_slot=0) - self._prefill_all_but_last_token(engine, kv_cache_manager, - resource_manager, context) - context.py_return_context_logits = True - batch = ScheduledRequests() - batch.context_requests_last_chunk = [context] - graph_runner.key = KeyType(batch_size=1, draft_len=0) - - self._forward(engine, kv_cache_manager, resource_manager, batch) - - lookup = graph_runner.maybe_get_cuda_graph.call_args - self.assertEqual(lookup.args[0].generation_requests, [context]) - self.assertEqual(lookup.kwargs["promoted_context_request_ids"], - frozenset({context.py_request_id})) - graph_runner.replay.assert_called_once() - self.assertEqual(graph_runner.replayed_inputs[0]["input_ids"].tolist(), - [44]) - - def test_promoted_context_uses_prompt_token_during_overlap(self) -> None: - engine, kv_cache_manager, resource_manager, graph_runner = \ - self._make_engine() - context = self._request([11, 22, 33, 44], 1, seq_slot=0) - generation = self._request([50, 51, 52, 53, 54], 2, seq_slot=1) - self._prefill_all_but_last_token(engine, kv_cache_manager, - resource_manager, context) - self._prefill(engine, kv_cache_manager, resource_manager, generation) - # The chunked prefill stands in for a reused prefix, so the promoted - # row is the context's first forward and latches its cached tokens. - context._cached_tokens_set = False - # A promoted context row must ignore any stale overlap slot. - context.py_batch_idx = 3 - context.py_num_compressed_tokens = 1 - generation.py_batch_idx = 1 - new_tokens = torch.zeros((1, 4, 1), dtype=torch.int32, device="cuda") - new_tokens[0, 0, 0] = 999 - new_tokens[0, 1, 0] = 777 - batch = ScheduledRequests() - batch.context_requests_last_chunk = [context] - batch.generation_requests = [generation] - graph_runner.key = KeyType(batch_size=2, draft_len=0) - - self._forward(engine, kv_cache_manager, resource_manager, batch, - SimpleNamespace(new_tokens=new_tokens)) - - replayed = graph_runner.replayed_inputs[0] - self.assertEqual(replayed["input_ids"].tolist(), [44, 777]) - self.assertEqual(replayed["position_ids"].flatten().tolist(), [3, 5]) - attn_metadata = engine.attn_metadata - self.assertEqual(attn_metadata.num_contexts, 0) - self.assertEqual( - attn_metadata.kv_cache_params.num_cached_tokens_per_seq, [2, 5]) - self.assertEqual(context.cached_tokens, 3) - - def test_steady_overlap_generation_advances_positions(self) -> None: - engine, kv_cache_manager, resource_manager, _ = self._make_engine() - requests = [ - self._request([10, 11, 12, 13], 1, seq_slot=0), - self._request([20, 21, 22, 23, 24, 25], 2, seq_slot=1), - ] - self._prefill(engine, kv_cache_manager, resource_manager, *requests) - for batch_idx, request in enumerate(requests): - request.py_batch_idx = batch_idx - batch = ScheduledRequests() - batch.generation_requests = requests - steps = ((100, 200), (101, 201), (102, 202)) - - for step, tokens in enumerate(steps): - new_tokens = torch.tensor(tokens, dtype=torch.int32, - device="cuda").view(1, 2, 1) - self._forward(engine, kv_cache_manager, resource_manager, batch, - SimpleNamespace(new_tokens=new_tokens)) - model_kwargs = engine.model.recorded_kwargs[-1] - with self.subTest(step=step): - self.assertEqual(model_kwargs["input_ids"].tolist(), - list(tokens)) - self.assertEqual( - model_kwargs["position_ids"].flatten().tolist(), - [4 + step, 6 + step]) - # The sampler appends the previous step's token after the next - # step has been scheduled. - for request, token in zip(requests, tokens): - request.add_new_token(token, 0) - - def _make_mrope_engine(self): - return self._make_engine( - DummyModel(torch.half, rope_scaling={"type": "mrope"})) - - def _mrope_generation(self, engine, kv_cache_manager, resource_manager, - num_tokens: int, req_id: int, seq_slot: int, - delta: int | None) -> LlmRequest: - """A generation request at position `num_tokens - 1`. - - It carries an MRoPE delta when `delta` is not None, and no - multimodal data at all otherwise (a text-only prompt). - """ - request = self._request(list(range(1, num_tokens)), req_id, seq_slot) - self._prefill(engine, kv_cache_manager, resource_manager, request) - request.add_new_token(num_tokens, 0) - if delta is None: - request.py_multimodal_data = {} - else: - request.py_multimodal_data = { - "mrope_config": { - "mrope_position_deltas": - torch.tensor([[delta]], dtype=torch.int32) - }, - } - return request - - def _forward_generation(self, engine, kv_cache_manager, resource_manager, - requests: list[LlmRequest]) -> dict: - batch = ScheduledRequests() - batch.generation_requests = requests - self._forward(engine, kv_cache_manager, resource_manager, batch) - return engine.model.recorded_kwargs[-1] - - def test_mrope_generation_with_partial_segments(self) -> None: - """A real multimodal span next to a padding dummy without MRoPE data.""" - engine, kv_cache_manager, resource_manager, _ = \ - self._make_mrope_engine() - multimodal = self._mrope_generation(engine, kv_cache_manager, - resource_manager, 4, 1, 0, 10) - multimodal.py_multimodal_data["multimodal_embedding"] = torch.ones( - (1, 1), dtype=torch.float16) - # A generation-only admission has no batch slot yet, which is when - # the context-phase payload is stripped. - multimodal.py_batch_idx = None - dummy = self._mrope_generation(engine, kv_cache_manager, - resource_manager, 6, 2, 1, None) - dummy.is_cuda_graph_dummy = True - - model_kwargs = self._forward_generation(engine, kv_cache_manager, - resource_manager, - [multimodal, dummy]) - - expected = torch.tensor([[[13, 5]]] * 3, - dtype=torch.int32, - device='cuda') - torch.testing.assert_close(model_kwargs["position_ids"], - expected, - atol=0, - rtol=0) - self.assertEqual( - model_kwargs["mrope_delta_write_seq_slots"].cpu().tolist(), [0]) - # Read slots are dense w.r.t. the generation batch: the padded dummy - # has no MRoPE metadata, so it resolves to the reserved zero slot - # (max_num_tokens * pp_size) rather than being dropped, which would - # shift every later request onto another request's delta. - self.assertEqual( - model_kwargs["mrope_delta_read_seq_slots"].cpu().tolist(), - [0, engine.max_num_tokens]) - self.assertNotIn("multimodal_embedding", multimodal.py_multimodal_data) - - def test_mrope_mixed_text_only_keeps_deltas_dense(self) -> None: - """A text-only request between two multimodal ones must not compact the - MRoPE delta read slots. - - The attention kernel indexes `mrope_position_deltas` by generation batch - index, so a list that skips the text-only request would hand request 2's - delta to the text-only request and read out of bounds for request 2. - """ - engine, kv_cache_manager, resource_manager, _ = \ - self._make_mrope_engine() - # (num_tokens, req_id, seq_slot, delta); the middle one is text-only. - requests = [ - self._mrope_generation(engine, kv_cache_manager, resource_manager, - *args) - for args in ((4, 1, 0, 10), (5, 2, 1, None), (6, 3, 2, 20)) - ] - - model_kwargs = self._forward_generation(engine, kv_cache_manager, - resource_manager, requests) - - # One entry per generation request, in batch order; the reserved zero - # slot stands in for the text-only request's zero delta. - self.assertEqual( - model_kwargs["mrope_delta_read_seq_slots"].cpu().tolist(), - [0, engine.max_num_tokens, 2]) - # Only the two multimodal requests seed the seq-slot delta cache. - self.assertEqual( - model_kwargs["mrope_delta_write_seq_slots"].cpu().tolist(), [0, 2]) - # The position is num_tokens - 1, offset by the request's delta; the - # text-only request keeps the plain scalar position on all 3 axes. - expected = torch.tensor([[[13, 4, 25]]] * 3, - dtype=torch.int32, - device='cuda') - torch.testing.assert_close(model_kwargs["position_ids"], - expected, - atol=0, - rtol=0) - - def test_mrope_multi_beam_generation_positions(self) -> None: - """Every beam row gets its request's `past + delta` on all 3 axes.""" - beam_width = 2 - engine, kv_cache_manager = create_model_engine_and_kvcache( - llm_args=TorchLlmArgs(model="dummy", - max_batch_size=4, - max_num_tokens=64, - max_beam_width=beam_width), - model=DummyModel(torch.half, rope_scaling={"type": "mrope"})) - self.addCleanup(kv_cache_manager.shutdown) - preparer = engine._runner._inputs - # (num_tokens, delta): one delta arrives on the host and one is - # already a CUDA int32 scalar. - cases = ((4, torch.tensor([[10]], dtype=torch.int32)), - (7, torch.tensor([[-3]], dtype=torch.int32, device="cuda"))) - requests = [] - for seq_slot, (num_tokens, delta) in enumerate(cases): - request = LlmRequest( - request_id=seq_slot + 1, - max_new_tokens=1, - input_tokens=list(range(1, num_tokens + 1)), - sampling_config=tensorrt_llm.bindings.SamplingConfig( - SamplingParams( - n=beam_width, - use_beam_search=True)._get_sampling_config()), - is_streaming=False) - request.py_seq_slot = seq_slot - request.py_multimodal_data = { - "mrope_config": { - "mrope_position_deltas": delta - } - } - requests.append(request) - host = HostInputs(use_mrope=True, - mrope_dummy_seq_slot=engine.max_num_tokens, - extra_inputs=preparer._new_extra_inputs(), - context_prompt_lookahead=None) - host.lists.generation_requests.extend(requests) - - batched_mrope_generation_updates: list[tuple[int, int, - torch.Tensor]] = [] - preparer._pack_generation_requests( - host, - frozenset(), - None, - batched_mrope_generation_updates=batched_mrope_generation_updates) - host.total_num_tokens = len(host.lists.position_ids) - preparer._copy_position_ids_to_device( - host, - batched_mrope_generation_updates=batched_mrope_generation_updates) - - # Both deltas take the batched scalar path, one row per beam. - self.assertEqual( - [row for row, _, _ in batched_mrope_generation_updates], - [0, 1, 2, 3]) - self.assertEqual(host.lists.mrope_position_ids, []) - # Without a previous tensor the past length is num_tokens - 1. - expected_row = [ - num_tokens - 1 + int(delta) for num_tokens, delta in cases - for _ in range(beam_width) - ] - self.assertEqual(expected_row, [13, 13, 3, 3]) - torch.testing.assert_close( - preparer._buffers.mrope_position_ids_cuda[:, :, :4], - torch.tensor([[expected_row]] * 3, dtype=torch.int32, - device="cuda"), - atol=0, - rtol=0) - - def test_mrope_mixed_batched_and_fallback_positions_match_fallback( - self) -> None: - """Mixed scalar batching matches the old per-request writes exactly.""" - engine, kv_cache_manager = create_model_engine_and_kvcache( - llm_args=TorchLlmArgs(model="dummy", - max_batch_size=2, - max_num_tokens=64), - model=DummyModel(torch.half, rope_scaling={"type": "mrope"})) - self.addCleanup(kv_cache_manager.shutdown) - preparer = engine._runner._inputs - sentinel = 123 - - def run( - delta_dtypes: tuple[torch.dtype, torch.dtype] - ) -> tuple[HostInputs, list[tuple[int, int, torch.Tensor]], - torch.Tensor]: - requests = [] - for seq_slot, (num_tokens, delta, dtype) in enumerate( - zip((4, 7), (10, -3), delta_dtypes)): - request = LlmRequest( - request_id=seq_slot + 1, - max_new_tokens=1, - input_tokens=list(range(1, num_tokens + 1)), - sampling_config=tensorrt_llm.bindings.SamplingConfig( - SamplingParams()._get_sampling_config()), - is_streaming=False) - request.py_seq_slot = seq_slot - request.py_multimodal_data = { - "mrope_config": { - "mrope_position_deltas": - torch.tensor([[delta]], dtype=dtype, device="cuda") - } - } - requests.append(request) - host = HostInputs(use_mrope=True, - mrope_dummy_seq_slot=engine.max_num_tokens, - extra_inputs=preparer._new_extra_inputs(), - context_prompt_lookahead=None) - host.lists.generation_requests.extend(requests) - batched_mrope_generation_updates: list[tuple[int, int, - torch.Tensor]] = [] - preparer._pack_generation_requests( - host, - frozenset(), - None, - batched_mrope_generation_updates=batched_mrope_generation_updates - ) - host.total_num_tokens = len(host.lists.position_ids) - preparer._buffers.mrope_position_ids_cuda.fill_(sentinel) - preparer._copy_position_ids_to_device( - host, - batched_mrope_generation_updates=batched_mrope_generation_updates - ) - positions = preparer._buffers.mrope_position_ids_cuda.clone() - return host, batched_mrope_generation_updates, positions - - mixed_host, mixed_updates, mixed = run((torch.int32, torch.int64)) - _, fallback_updates, fallback = run((torch.int64, torch.int64)) - - self.assertEqual([row for row, _, _ in mixed_updates], [0]) - self.assertEqual( - [(start, end) - for start, end, _ in mixed_host.lists.mrope_position_ids], - [(1, 2)]) - self.assertEqual(fallback_updates, []) - torch.testing.assert_close(mixed, fallback, atol=0, rtol=0) - self.assertTrue(torch.all(mixed[:, :, 2:] == sentinel).item()) - - def test_mrope_all_text_only_drops_deltas(self) -> None: - """A generation batch with no MRoPE metadata at all emits no delta - tensors, so the steady-state generation fast path stays reachable.""" - engine, kv_cache_manager, resource_manager, _ = \ - self._make_mrope_engine() - requests = [ - self._mrope_generation(engine, kv_cache_manager, resource_manager, - 4, 1, 0, None), - self._mrope_generation(engine, kv_cache_manager, resource_manager, - 6, 2, 1, None), - ] - - model_kwargs = self._forward_generation(engine, kv_cache_manager, - resource_manager, requests) - - # An all-zero delta vector is identical to passing no deltas at all. - self.assertNotIn("mrope_delta_read_seq_slots", model_kwargs) - self.assertNotIn("mrope_delta_write_seq_slots", model_kwargs) - # MRoPE models keep the (3,1,N) layout even with no mrope work. - expected = torch.tensor([[[3, 5]]] * 3, - dtype=torch.int32, - device='cuda') - torch.testing.assert_close(model_kwargs["position_ids"], - expected, - atol=0, - rtol=0) - - def test_promoted_mrope_context_uses_decode_state_contract(self) -> None: - engine, kv_cache_manager, resource_manager, graph_runner = \ - self._make_mrope_engine() - engine.model.config.vocab_size = 100 - context = self._request([11, 22, 33, 44], 1, seq_slot=0) - self._prefill_all_but_last_token(engine, kv_cache_manager, - resource_manager, context) - mrope_delta = torch.tensor([[10]], dtype=torch.int32) - context.py_mrope_position_delta = mrope_delta - context.py_mrope_delta_cache_slot = context.py_seq_slot - context.py_multimodal_data = { - "mrope_config": { - "mrope_position_deltas": mrope_delta, - }, - "multimodal_embedding": torch.ones((1, 1), dtype=torch.float16), - } - batch = ScheduledRequests() - batch.context_requests_last_chunk = [context] - graph_runner.key = KeyType(batch_size=1, draft_len=0) - - self._forward(engine, kv_cache_manager, resource_manager, batch) - - self.assertEqual( - graph_runner.maybe_get_cuda_graph.call_args. - kwargs["promoted_context_request_ids"], - frozenset({context.py_request_id})) - replayed = graph_runner.replayed_inputs[0] - self.assertEqual(replayed["input_ids"].tolist(), [44]) - expected_positions = torch.full((3, 1, 1), - 13, - dtype=torch.int32, - device="cuda") - torch.testing.assert_close(replayed["position_ids"], - expected_positions, - atol=0, - rtol=0) - self.assertEqual(replayed["mrope_delta_read_seq_slots"].cpu().tolist(), - [0]) - self.assertNotIn("mrope_delta_write_seq_slots", replayed) - attn_metadata = engine.attn_metadata - self.assertEqual( - attn_metadata.kv_cache_params.num_cached_tokens_per_seq, [3]) - self.assertEqual(attn_metadata.num_contexts, 0) - - -class PyTorchModelEngineTestCase(unittest.TestCase): - - def test_warmup_builds_dsa_attn_metadata_when_no_forward_ran(self) -> None: - """The DSA top-K pre-compile hooks read attn_metadata; when every - warmup forward was skipped it is built on demand so the engines still - compile before serving.""" - from tensorrt_llm._torch.attention.backends.sparse.dsa import \ - DSAtrtllmAttentionMetadata - - engine = object.__new__(DecoderWarmup) - engine._config = SimpleNamespace( - attention_backend=SimpleNamespace( - Metadata=DSAtrtllmAttentionMetadata), - original_max_draft_len=2, - cuda_graph_batch_sizes=[8], - spec_config=None, - ) - engine._metadata = DecoderMetadata(config=engine._config, - mapping=None, - model=None, - buffers=None, - attention_runtime_features=None) - metadata = Mock(spec=DSAtrtllmAttentionMetadata) - - def build(kv_cache_manager, draft_kv_cache_manager): - engine._metadata.attn_metadata = metadata - return metadata - - engine._metadata.setup_attention = Mock(side_effect=build) - kv_cache_manager = Mock() - resource_manager = Mock() - resource_manager.get_resource_manager.return_value = kv_cache_manager - - engine._ensure_dsa_attn_metadata_for_warmup(resource_manager) - engine._warmup_cute_dsl_radix_topk() - - engine._metadata.setup_attention.assert_called_once_with( - kv_cache_manager, None) - metadata.warmup_cute_dsl_radix_topk.assert_called_once_with(3) - metadata.warmup_selfsampling_topk.assert_called_once_with( - 3, batch_sizes=[8]) - - def test_warmup_does_not_build_metadata_for_non_dsa_backend(self) -> None: - from tensorrt_llm._torch.attention.backends.trtllm import \ - TrtllmAttentionMetadata - - engine = object.__new__(DecoderWarmup) - engine._config = SimpleNamespace(attention_backend=SimpleNamespace( - Metadata=TrtllmAttentionMetadata)) - engine._metadata = DecoderMetadata(config=engine._config, - mapping=None, - model=None, - buffers=None, - attention_runtime_features=None) - engine._metadata.setup_attention = Mock() - - engine._ensure_dsa_attn_metadata_for_warmup(Mock()) - engine._warmup_cute_dsl_radix_topk() - - engine._metadata.setup_attention.assert_not_called() - - def test_breakable_rejects_multimodal_models(self) -> None: - engine = object.__new__(ModelEngine) - engine.model = DummyLegacyMultimodalIndexModel() - engine.input_processor = None - engine.llm_args = SimpleNamespace( - prefill_cuda_graph_backend=PrefillCudaGraphBackend.BREAKABLE, - disable_mm_encoder=False) - - self.assertTrue(is_multimodal(engine.model, engine.input_processor)) - with self.assertRaisesRegex(ValueError, "multimodal models"): - engine._validate_breakable_cuda_graph_compatibility() - - def test_breakable_allows_text_decoder_with_multimodal_processor( - self) -> None: - engine = object.__new__(ModelEngine) - engine.model = Mock(spec=DecoderModelForCausalLM) - engine.input_processor = Mock(spec=BaseMultimodalInputProcessor) - engine.llm_args = SimpleNamespace( - prefill_cuda_graph_backend=PrefillCudaGraphBackend.BREAKABLE, - disable_mm_encoder=False) - - self.assertTrue(is_multimodal(engine.model, engine.input_processor)) - engine._validate_breakable_cuda_graph_compatibility() - - def test_breakable_allows_multimodal_wrapper_in_text_only_mode( - self) -> None: - engine = object.__new__(ModelEngine) - engine.model = DummyLegacyMultimodalIndexModel() - engine.model.llm = Mock(spec=DecoderModelForCausalLM) - engine.model.mm_encoder = None - engine.input_processor = Mock(spec=BaseMultimodalInputProcessor) - engine.llm_args = SimpleNamespace( - prefill_cuda_graph_backend=PrefillCudaGraphBackend.BREAKABLE, - disable_mm_encoder=True) - - self.assertTrue(is_multimodal(engine.model, engine.input_processor)) - engine._validate_breakable_cuda_graph_compatibility() - - def test_restore_after_refit_rewraps_without_recapture(self) -> None: - """Restore re-wraps torch.compile and does nothing else. - - Refit moves no tensor, so the piecewise captures stay valid across it; - recapturing them costs a full re-capture per refit and leaves PWCG cold - in between. This pins the cheap path -- reintroducing the recapture, - the torch.compiler.reset() or the gc.collect() fails here. - - Replaces test_piecewise_refit_recapture_preserves_padding_dummy_requests, - which asserted the refit path recaptures. That path no longer exists. - """ - cuda_graph_runner = SimpleNamespace( - padding_dummy_requests={0: object()}, - allow_capture=lambda: nullcontext(), - ) - model_engine = SimpleNamespace( - warmup_scope=nullcontext, - _torch_compile_backend=object(), - llm_args=SimpleNamespace(torch_compile_config=SimpleNamespace( - enable_fullgraph=False)), - cuda_graph_runner=cuda_graph_runner, - _apply_torch_compile=Mock(), - _capture_piecewise_cuda_graphs=Mock(), - ) - - with patch("torch.compiler.reset") as reset, \ - patch("gc.collect") as collect: - ModelEngine.restore_compiled_model_after_refit( - model_engine, resource_manager=object()) - - model_engine._apply_torch_compile.assert_called_once() - model_engine._capture_piecewise_cuda_graphs.assert_not_called() - reset.assert_not_called() - # No full collection either: it ran for seconds inside the refit - # bubble and the replaced wrapper is freed by reference count. - collect.assert_not_called() - - def test_restore_after_refit_noop_when_torch_compile_disabled(self) -> None: - """An eager engine must not be silently handed a compiled wrapper.""" - model_engine = SimpleNamespace( - _torch_compile_backend=None, - _apply_torch_compile=Mock(), - ) - ModelEngine.restore_compiled_model_after_refit( - model_engine, resource_manager=object()) - model_engine._apply_torch_compile.assert_not_called() - - def test_restore_after_refit_never_touches_the_kv_cache(self) -> None: - """The re-wrap hook must not run the warm-up KV-cache cleanup. - - ``warmup_with_kv_cache_cleanup`` ends with - ``check_invalid_values_in_kv_cache(fill_with_zero=True)``, which - zero-fills every attention KV pool whether or not it found a NaN. That - is right after a capture warm-up that ran forwards on placeholder - inputs; restore never runs a forward, and under an in-flight weight - update the zero-fill would wipe the KV of the requests kept across the - refit. So the fully decorated hook must never even look up the KV cache - manager -- with torch.compile on or off. - """ - for compile_enabled in (True, False): - with self.subTest(compile_enabled=compile_enabled): - resource_manager = Mock() - model_engine = SimpleNamespace( - warmup_scope=nullcontext, - _torch_compile_backend=Mock() if compile_enabled else None, - llm_args=SimpleNamespace( - torch_compile_config=SimpleNamespace( - enable_fullgraph=False)), - _apply_torch_compile=Mock(), - ) - with patch("gc.collect"): - ModelEngine.restore_compiled_model_after_refit( - model_engine, resource_manager) - resource_manager.get_resource_manager.assert_not_called() - self.assertEqual(model_engine._apply_torch_compile.call_count, - 1 if compile_enabled else 0) - - def test_unwrap_for_refit_gates_on_torch_compile_not_piecewise( - self) -> None: - """Compile-on/piecewise-off must still unwrap. - - This is the original defect: the hook was gated on the piecewise flag, - so TorchCompileConfig() -- compile on, piecewise off -- got no refit - lifecycle and refit loaded into a wrapped model, matching nothing. - """ - for piecewise in (True, False): - with self.subTest(piecewise=piecewise): - model_engine = SimpleNamespace( - llm_args=SimpleNamespace(prefill_cuda_graph_backend=( - PrefillCudaGraphBackend.PIECEWISE if piecewise else None - )), - _torch_compile_backend=Mock(), - _remove_torch_compile=Mock(), - ) - unwrap_core = (ModelEngine.unwrap_compiled_model_for_refit) - with patch("torch.cuda.synchronize"): - unwrap_core(model_engine) - model_engine._remove_torch_compile.assert_called_once() - # The captures are kept, not cleared: refit moves no tensor, so - # clearing them only forces PWCG cold for no benefit. - model_engine._torch_compile_backend \ - .clear_piecewise_cuda_graphs.assert_not_called() - - # Eager engine: the hook must be a no-op. - eager = SimpleNamespace( - _torch_compile_backend=None, - _remove_torch_compile=Mock(), - ) - ModelEngine.unwrap_compiled_model_for_refit(eager) - eager._remove_torch_compile.assert_not_called() - - def test_refit_lifecycle_aliases_point_at_renamed_hooks(self) -> None: - """NeMo-RL resolves these with getattr() and falls back to doing nothing. - - Dropping an alias would therefore not raise -- it would silently skip - the unwrap and reintroduce the refit corruption. - """ - self.assertIs(ModelEngine.release_piecewise_cuda_graphs_for_refit, - ModelEngine.unwrap_compiled_model_for_refit) - self.assertIs(ModelEngine.recapture_piecewise_cuda_graphs_after_refit, - ModelEngine.restore_compiled_model_after_refit) - - def test_build_request_multimodal_input_skips_when_cache_disabled( - self) -> None: - request = LlmRequest( - request_id=1, - max_new_tokens=1, - input_tokens=[0, 1, 2], - sampling_config=tensorrt_llm.bindings.SamplingConfig(1), - is_streaming=False, - multimodal_hashes=[[1, 2, 3, 4, 5, 6, 7, 8]], - multimodal_positions=[1], - multimodal_lengths=[1], - multimodal_uuids=["image-0"], - ) - - # With the encoder cache disabled, nothing consumes `multimodal_input`, - # so it should not be built at all. - self.assertIsNone( - _build_request_multimodal_input(request, cache_enabled=False)) - - def test_spec_decode_capture_restores_kv_lens_between_warmups(self) -> None: - attn_metadata = Mock() - attn_metadata.num_seqs = 1 - attn_metadata.kv_lens_cuda = torch.tensor([4095], dtype=torch.int32) - - saved_kv_lens_cuda = _save_spec_decode_capture_state( - attn_metadata, enable_spec_decode=True) - - # CUDA graph capture performs two eager warmup forwards. A speculative - # draft loop may advance the static attention metadata during each - # forward, but the next warmup must start from the original input. - for _ in range(2): - attn_metadata.kv_lens_cuda.add_(1) - _restore_spec_decode_capture_state(attn_metadata, - saved_kv_lens_cuda) - self.assertEqual(attn_metadata.kv_lens_cuda.tolist(), [4095]) - - self.assertEqual(attn_metadata.on_update_kv_lens.call_count, 2) - - def test_decoder_buffers_are_allocated_only_for_decoder_runner( - self) -> None: - llm_args = TorchLlmArgs( - model="dummy", - max_batch_size=4, - max_num_tokens=64, - max_beam_width=2, - cuda_graph_config=CudaGraphConfig(batch_sizes=[1, 2, 4])) - for is_generation in (True, False): - with self.subTest(is_generation=is_generation): - engine = DummyModelEngine(llm_args, - torch.half, - is_generation=is_generation) - - runner = engine._runner - if is_generation: - self.assertIsInstance(runner, DecoderRunner) - self.assertIsNotNone(engine.cuda_graph_runner) - self.assertEqual( - tuple(runner._buffers.cache_indirection.shape), - (4, 2, engine.max_seq_len)) - else: - self.assertIsInstance(runner, PoolingRunner) - self.assertIsNone(engine.cuda_graph_runner) - self.assertEqual( - isinstance(runner._buffers, DecoderInputBuffers), - is_generation) - - def test_no_kv_cache_runners_reject_speculative_config(self) -> None: - for mm_encoder_only, runner_name in ((False, "PoolingRunner"), - (True, "MultimodalEncoderRunner")): - with self.subTest(runner=runner_name): - llm_args = TorchLlmArgs(model="dummy", - max_batch_size=4, - max_num_tokens=64, - mm_encoder_only=mm_encoder_only) - with self.assertRaisesRegex( - ValueError, - f"{runner_name} does not support speculative decoding"): - DummyModelEngine( - llm_args, - torch.half, - spec_config=SADecodingConfig(max_draft_len=1), - is_generation=False) - - def test_runner_engine_shared_entries_skip_decoder_state(self) -> None: - runner = Mock() - engine = object.__new__(ModelEngine) - engine._runner = runner - engine._torch_compile_backend = None - - engine.wait_for_input_copy() - engine.release_executor_resources(release_model_owned_graphs=False) - engine.init_cuda_graph_lora_manager(Mock()) - - runner.wait_for_input_copy.assert_called_once_with() - runner.release_graphs.assert_called_once_with( - release_nccl_window_owners=True) - runner.init_cuda_graph_lora_manager.assert_not_called() - - def test_decoder_engine_shared_entries_use_decoder_state(self) -> None: - engine = object.__new__(ModelEngine) - decoder = object.__new__(DecoderRunner) - engine._runner = decoder - engine._torch_compile_backend = None - decoder._inputs = object.__new__(DecoderInputPreparer) - decoder._inputs._prepare_inputs_event = Mock() - decoder._graphs = DecoderCudaGraphs( - cuda_graph_runner=Mock(), - breakable_cuda_graph_runner=Mock(), - config=None, - mapping=Mapping(), - model=None, - sparse_attention_config=None, - metrics={}, - dummy_requests=None, - agreement=None, - metadata=None, - get_runtime_tokens_per_gen_step=None, - dynamic_draft_len_mapping=None, - lora=None, - forward_warmup=None, - torch_compile_bypass=nullcontext, - context_promotion_supported=True, - ) - decoder._metadata = SimpleNamespace(attn_metadata=Mock()) - - engine.wait_for_input_copy() - engine.release_executor_resources(release_model_owned_graphs=False) - - decoder._inputs._prepare_inputs_event.synchronize.assert_called_once_with( - ) - decoder.cuda_graph_runner.clear.assert_called_once_with( - release_nccl_window_owners=True) - decoder._graphs.breakable_cuda_graph_runner.clear.assert_called_once_with( - release_nccl_window_owners=True) - self.assertIsNone(decoder.attn_metadata) - - def test_prepare_inputs_returns_steady_gen_fast_inputs(self) -> None: - inputs = _input_preparer() - inputs._config = SimpleNamespace(spec_config=None) - inputs.guided_decoder = Mock() - inputs._can_use_steady_gen_fast_prepare = Mock(return_value=True) - fast_inputs = PreparedInputs({"prepared": True}) - inputs._apply_steady_gen_fast_prepare = Mock(return_value=fast_inputs) - batch = ScheduledRequests() - new_tokens = torch.zeros(1) - - result = inputs.prepare_inputs( - batch, - Mock(), - Mock(), - new_tensors_device=SimpleNamespace(new_tokens=new_tokens), - enable_spec_decode=False, - runtime_draft_len=3, - is_dummy=True) - - self.assertIs(result.kwargs, fast_inputs.kwargs) - self.assertEqual(result.runtime_draft_len, 3) - inputs.guided_decoder.add_batch.assert_called_once_with( - batch, new_tokens=new_tokens, runtime_draft_len=3) - inputs._can_use_steady_gen_fast_prepare.assert_called_once_with( - batch, new_tokens, None, None, True) - - def test_prepare_inputs_reports_tree_draft_length(self) -> None: - inputs = _input_preparer() - inputs._config = SimpleNamespace(spec_config=SimpleNamespace( - is_linear_tree=False, tokens_per_gen_step=8)) - inputs._can_use_steady_gen_fast_prepare = Mock(return_value=True) - inputs._apply_steady_gen_fast_prepare = Mock( - return_value=PreparedInputs({})) - - for enable_spec_decode, expected in ((True, 7), (False, 2)): - with self.subTest(enable_spec_decode=enable_spec_decode): - result = inputs.prepare_inputs( - ScheduledRequests(), - Mock(), - Mock(), - enable_spec_decode=enable_spec_decode, - runtime_draft_len=2, - is_dummy=True) - self.assertEqual(result.runtime_draft_len, expected) - - def test_steady_gen_fast_prepare_skips_dummy_passes(self) -> None: - inputs = _input_preparer() - request = _make_request_stub(7) - inputs._steady_gen_cache = { - "num_requests": 1, - "request_ids": [request.py_request_id], - } - batch = ScheduledRequests() - batch.generation_requests = [request] - new_tokens = torch.zeros(1) - - self.assertTrue( - inputs._can_use_steady_gen_fast_prepare(batch, new_tokens, None, - None, False)) - self.assertFalse( - inputs._can_use_steady_gen_fast_prepare(batch, new_tokens, None, - None, True)) - - def test_padding_request_input_id_rows_are_zeroed(self) -> None: - model_engine, kv_cache_manager = create_model_engine_and_kvcache(self) - resource_manager = ResourceManager( - {ResourceManagerType.KV_CACHE_MANAGER: kv_cache_manager}) - attn_metadata = AttentionMetadata(max_num_requests=4, - max_num_tokens=32, - kv_cache_manager=kv_cache_manager) - attn_metadata.is_cuda_graph = False - - generation = _create_request_with_tokens([50, 51, 52, 53, 54], 1) - generation.py_seq_slot = 0 - generation.py_batch_idx = 0 - padding = _create_request_with_tokens([60, 61, 62], 2) - padding.is_cuda_graph_dummy = True - padding.py_seq_slot = 1 - - graph_batch = ScheduledRequests() - graph_batch.generation_requests = [generation, padding] - new_tokens = torch.zeros((1, 4, 1), dtype=torch.int32, device="cuda") - new_tokens[0, 0, 0] = 777 - overlap_state = SimpleNamespace(new_tokens=new_tokens) - inputs = model_engine._runner._inputs - inputs._can_use_steady_gen_fast_prepare = Mock(return_value=False) - # The padding row is written by neither input-id copy, so poison the - # buffer with an out-of-vocabulary id to catch a stale read. - model_engine._runner._buffers.input_ids_cuda.fill_(1 << 30) - - prepared = inputs.prepare_inputs( - scheduled_requests=graph_batch, - kv_cache_manager=kv_cache_manager, - attn_metadata=attn_metadata, - new_tensors_device=overlap_state, - resource_manager=resource_manager, - enable_spec_decode=False, - runtime_draft_len=0, - is_dummy=False, - ) - - self.assertEqual(prepared.kwargs["input_ids"][:2].cpu().tolist(), - [777, 0]) - kv_cache_manager.shutdown() - - def test_promoted_context_precedes_speculative_overlap_generation( - self) -> None: - spec_config = SADecodingConfig( - max_draft_len=1, - draft_len_schedule={1: 1}, - ) - model_engine, kv_cache_manager = create_model_engine_and_kvcache( - self, spec_config=spec_config) - resource_manager = ResourceManager( - {ResourceManagerType.KV_CACHE_MANAGER: kv_cache_manager}) - attn_metadata = AttentionMetadata(max_num_requests=4, - max_num_tokens=32, - kv_cache_manager=kv_cache_manager) - attn_metadata.is_cuda_graph = False - # A bare Mock auto-vivifies every attribute, so optional metadata - # fields read by prepare_inputs must be pinned off explicitly. - spec_metadata = Mock( - _force_non_greedy_for_capture=False, - context_prompt_lookahead_tokens=None, - ) - - context = _create_request_with_tokens([11, 22, 33, 44], 1) - context.context_current_position = 3 - context.context_chunk_size = 1 - context.py_seq_slot = 0 - # A promoted context row must ignore any stale overlap slot. - context.py_batch_idx = 3 - - generation = _create_request_with_tokens([50, 51, 52, 53, 54], 2) - generation.py_seq_slot = 1 - generation.py_batch_idx = 1 - generation.py_needs_onehot_draft_probs = True - - graph_batch = ScheduledRequests() - graph_batch.generation_requests = [context, generation] - new_tokens = torch.zeros((1, 4, 1), dtype=torch.int32, device="cuda") - new_tokens[0, 0, 0] = 999 - new_tokens[0, 1, 0] = 777 - overlap_state = SampleStateTensorsSpec( - new_tokens=new_tokens, - new_tokens_lens=torch.ones(4, dtype=torch.int32, device="cuda"), - next_draft_tokens=torch.zeros((4, 1), - dtype=torch.int32, - device="cuda"), - ) - - prepared = model_engine._runner._inputs.prepare_inputs( - scheduled_requests=graph_batch, - kv_cache_manager=kv_cache_manager, - attn_metadata=attn_metadata, - spec_metadata=spec_metadata, - new_tensors_device=overlap_state, - resource_manager=resource_manager, - promoted_context_request_ids=frozenset({context.py_request_id}), - enable_spec_decode=True, - runtime_draft_len=0, - is_dummy=False, - ) - - inputs = prepared.kwargs - self.assertEqual(inputs["input_ids"][:2].cpu().tolist(), [44, 777]) - self.assertEqual(inputs["position_ids"][0, :2].cpu().tolist(), [3, 4]) - self.assertEqual(attn_metadata.request_ids, - [context.py_request_id, generation.py_request_id]) - self.assertEqual( - attn_metadata.kv_cache_params.num_cached_tokens_per_seq, [3, 5]) - self.assertEqual( - model_engine._runner._buffers.previous_batch_indices_cuda[:1].cpu(). - tolist(), [1]) - self.assertEqual( - model_engine._runner._buffers.previous_pos_id_offsets_cuda[:2].cpu( - ).tolist(), [0, 1]) - self.assertEqual(attn_metadata.num_contexts, 0) - self.assertEqual(spec_metadata.request_ids, - [context.py_request_id, generation.py_request_id]) - self.assertFalse(generation.py_needs_onehot_draft_probs) - spec_metadata.write_padding_onehot_draft_probs.assert_called_once_with( - [generation.py_seq_slot], 0) - kv_cache_manager.shutdown() - - def test_spec_decode_graph_step_gather_kernel_matches_torch_path( - self) -> None: - """A CUDA graph decode step of a speculative engine whose overlap - gathers run as one StepInputGather launch writes what the torch path - writes.""" - from tensorrt_llm._torch.attention.backends.trtllm import \ - TrtllmAttentionMetadata - from tensorrt_llm._torch.kernels.spec_step_copies import \ - op as spec_step_copies - if not spec_step_copies.is_supported(): - self.skipTest("the step-copy kernels run on SM 100") - max_draft_len = 3 - tokens_per_step = max_draft_len + 1 - model_engine, kv_cache_manager = create_model_engine_and_kvcache( - self, spec_config=SADecodingConfig(max_draft_len=max_draft_len)) - inputs = model_engine._runner._inputs - buffers = model_engine._runner._buffers - step_gather = inputs._step_input_gather - self.assertIsNotNone(step_gather) - resource_manager = ResourceManager( - {ResourceManagerType.KV_CACHE_MANAGER: kv_cache_manager}) - - # Three requests that ran in the previous step, in slots out of order. - slots = [5, 2, 7] - requests = kv_cache_manager.add_dummy_requests( - [1, 2, 3], - token_nums=[10, 17, 33], - is_gen=True, - max_num_draft_tokens=max_draft_len) - for request, slot in zip(requests, slots): - request.is_dummy_request = False - request.py_seq_slot = slot - batch = ScheduledRequests() - batch.generation_requests = requests - attn_metadata = model_engine._runner._metadata.setup_attention( - kv_cache_manager) - self.assertIs(type(attn_metadata), TrtllmAttentionMetadata) - graph_metadata = attn_metadata.create_cuda_graph_metadata( - len(requests), False, max_draft_len) - spec_metadata = Mock( - _force_non_greedy_for_capture=False, - context_prompt_lookahead_tokens=None, - ) - - # The sampler's slot stores as the previous step left them. - num_slots = 8 - generator = torch.Generator(device="cuda").manual_seed(0) - previous = SampleStateTensorsSpec( - new_tokens=torch.randint(0, - 1000, (tokens_per_step, num_slots, 1), - generator=generator, - dtype=torch.int32, - device="cuda"), - new_tokens_lens=torch.randint(1, - tokens_per_step + 1, (num_slots, ), - generator=generator, - dtype=torch.int32, - device="cuda"), - next_draft_tokens=torch.randint(0, - 1000, (num_slots, max_draft_len), - generator=generator, - dtype=torch.int32, - device="cuda"), - ) - - def step_inputs(step_input_gather): - """The buffers one step writes, from a filler they all start at.""" - inputs._step_input_gather = step_input_gather - for request, slot in zip(requests, slots): - request.py_batch_idx = slot - written = (buffers.input_ids_cuda, buffers.position_ids_cuda, - buffers.draft_tokens_cuda, - buffers.previous_pos_id_offsets_cuda, - buffers.previous_kv_lens_offsets_cuda) - for buffer in written: - buffer.fill_(-3) - inputs.prepare_inputs(scheduled_requests=batch, - kv_cache_manager=kv_cache_manager, - attn_metadata=graph_metadata, - spec_metadata=spec_metadata, - new_tensors_device=previous, - resource_manager=resource_manager, - enable_spec_decode=True, - runtime_draft_len=max_draft_len, - is_dummy=False) - torch.cuda.synchronize() - return [buffer.clone() for buffer in written] - - torch_path = step_inputs(None) - launched = [] - launch = step_gather.gather - - def recorded_launch(*args, **kwargs): - launched.append(launch(*args, **kwargs)) - return launched[-1] - - with patch.object(step_gather, "gather", side_effect=recorded_launch): - kernel_path = step_inputs(step_gather) - self.assertEqual(launched, [True]) - for expected, actual in zip(torch_path, kernel_path): - torch.testing.assert_close(actual, expected, rtol=0, atol=0) - # The previous rows' input ids came from the stores, by slot. - num_tokens = len(requests) * tokens_per_step - expected_input_ids = previous.new_tokens[:, slots, 0].t().reshape(-1) - self.assertEqual(kernel_path[0][:num_tokens].tolist(), - expected_input_ids.tolist()) - kv_cache_manager.shutdown() - - def test_multimodal_encoder_max_seq_len(self) -> None: - - class CapturingEncoder(torch.nn.Module, MultimodalEncoderMixin): - - def __init__(self) -> None: - super().__init__() - self.setup_args = None - self.max_seq_len = None - - def setup_attn_metadata(self, max_num_tokens: int) -> None: - self.setup_args = max_num_tokens - - def set_attn_max_seq_len(self, max_seq_len: int) -> None: - self.max_seq_len = max_seq_len - - encoder_max_num_tokens = 16384 - processor_max_num_tokens = 65536 - cases = [ - ({ - "image": processor_max_num_tokens - }, processor_max_num_tokens), - ({ - "image": 4096 - }, encoder_max_num_tokens), - ({}, encoder_max_num_tokens), - (None, encoder_max_num_tokens), - ] - for max_tokens_per_item, expected_max_seq_len in cases: - with self.subTest(max_tokens_per_item=max_tokens_per_item): - encoder = CapturingEncoder() - if max_tokens_per_item is None: - input_processor = Mock() - else: - input_processor = Mock( - spec=BaseMultimodalDummyInputsBuilder) - input_processor.get_mm_max_tokens_per_item.return_value = max_tokens_per_item - - setup_mm_encoder_attn_metadata(torch.nn.Sequential(encoder), - input_processor, - encoder_max_num_tokens, None) - - self.assertEqual(encoder.setup_args, encoder_max_num_tokens) - self.assertEqual(encoder.max_seq_len, expected_max_seq_len) - - def test_pad_generation_requests(self) -> None: - model_engine, kv_cache_manager = create_model_engine_and_kvcache(self) - resource_manager = ResourceManager( - {ResourceManagerType.KV_CACHE_MANAGER: kv_cache_manager}) - - def free_blocks() -> int: - return kv_cache_manager.get_kv_cache_stats().free_num_blocks - - batch_sizes_and_seqlens = [ - (5, 1), - (13, 1), - (5, 25), - ] - for (batch_size, max_seq_len) in batch_sizes_and_seqlens: - requests = [ - _create_request(max_seq_len, i) for i in range(batch_size) - ] - batch = ScheduledRequests() - batch.context_requests_last_chunk = requests - - pages_before = free_blocks() - with model_engine.cuda_graph_runner.pad_batch( - batch, resource_manager) as padded_batch: - # No padding for prefill - self.assertIs(batch, padded_batch) - self.assertEqual(free_blocks(), pages_before) - - batch = ScheduledRequests() - batch.generation_requests = requests - pages_before = free_blocks() - new_dummy_block = 1 if not model_engine.cuda_graph_runner.padding_dummy_requests else 0 - with model_engine.cuda_graph_runner.pad_batch( - batch, resource_manager) as padded_batch: - if batch_size < 8 and max_seq_len < 25: - self.assertEqual( - len(padded_batch.generation_requests) % 8, 0) - else: - # No padding if it would create too many concurrent requests. - # This requirement is not strictly required, but we should probably - # respect the requirement? - # The seqlen check makes sure we don't exceed the KV cache memory - # budget. - self.assertIs(batch, padded_batch) - self.assertEqual(free_blocks() + new_dummy_block, pages_before) - - kv_cache_manager.shutdown() - - def test_pad_batch_strips_cudagraph_dummies_on_clean_exit(self) -> None: - # Regression guard for the invariant that CUDAGraphRunner.pad_batch's - # `finally` strips every is_cuda_graph_dummy=True entry from - # scheduled_requests.generation_requests before the `with` block - # exits. Downstream consumers of scheduled_batch.generation_requests - # — including the per-iteration stats populate block in - # PyExecutor._update_iter_stats — rely on never observing - # cudagraph dummies. - model_engine, kv_cache_manager = create_model_engine_and_kvcache(self) - resource_manager = ResourceManager( - {ResourceManagerType.KV_CACHE_MANAGER: kv_cache_manager}) - - # batch_size=5 rounds up to 8 (nearest captured graph size in the - # fixture config) -> padding_size=3, deterministically. - real_batch_size = 5 - max_seq_len = 1 - real_requests = [ - _create_request(max_seq_len, i) for i in range(real_batch_size) - ] - real_ids = [id(r) for r in real_requests] - - batch = ScheduledRequests() - batch.generation_requests = list(real_requests) - - with model_engine.cuda_graph_runner.pad_batch( - batch, resource_manager) as padded_batch: - # Positive assertion that padding actually fired — guards - # against a vacuous pass where padding was a no-op. - self.assertGreater( - len(padded_batch.generation_requests), real_batch_size, - "padding did not fire; fixture config may have drifted " - "so that 5 no longer rounds up to 8") - # Every appended entry past the original count is a - # cudagraph-flagged dummy. - for req in padded_batch.generation_requests[real_batch_size:]: - self.assertTrue( - getattr(req, "is_cuda_graph_dummy", False), - "pad_batch appended a request without " - "is_cuda_graph_dummy=True") - # Real requests' identities and order are untouched. - self.assertEqual([ - id(r) - for r in padded_batch.generation_requests[:real_batch_size] - ], real_ids) - - # After the with-block: finally must have sliced off the padding. - self.assertEqual( - len(batch.generation_requests), real_batch_size, - "pad_batch.finally did not strip cudagraph dummies — " - "downstream consumers of scheduled_batch.generation_requests " - "would observe the leaked dummies") - for req in batch.generation_requests: - self.assertFalse( - getattr(req, "is_cuda_graph_dummy", False), - "cudagraph dummy leaked out of pad_batch's finally") - - kv_cache_manager.shutdown() - - def test_pad_batch_strips_cudagraph_dummies_on_exception(self) -> None: - # The strip must fire even when the body raises. This is the - # critical property of `finally` vs. a plain trailing statement — - # it guards the invariant on the error path. A refactor that - # accidentally dropped the `finally` would be caught here but not - # by the clean-exit variant. - model_engine, kv_cache_manager = create_model_engine_and_kvcache(self) - resource_manager = ResourceManager( - {ResourceManagerType.KV_CACHE_MANAGER: kv_cache_manager}) - - real_batch_size = 5 - real_requests = [_create_request(1, i) for i in range(real_batch_size)] - - batch = ScheduledRequests() - batch.generation_requests = list(real_requests) - - class _ForwardBoom(Exception): - pass - - with self.assertRaises(_ForwardBoom): - with model_engine.cuda_graph_runner.pad_batch( - batch, resource_manager) as padded_batch: - self.assertGreater(len(padded_batch.generation_requests), - real_batch_size) - raise _ForwardBoom() - - self.assertEqual(len(batch.generation_requests), real_batch_size) - for req in batch.generation_requests: - self.assertFalse(getattr(req, "is_cuda_graph_dummy", False)) - - kv_cache_manager.shutdown() - - def test_position_id_preparation(self) -> None: - model_engine, kv_cache_manager = create_model_engine_and_kvcache(self) - resource_manager = ResourceManager( - {ResourceManagerType.KV_CACHE_MANAGER: kv_cache_manager}) - - prompt_len = 256 - requests = [_create_request(prompt_len, 0)] - - # Prefill run - batch = ScheduledRequests() - batch.context_requests_last_chunk = requests - _prepare_kv_cache_resources(kv_cache_manager, batch) - model_engine.forward(batch, resource_manager) - - expected_prefill_pos_ids = torch.arange(0, - prompt_len, - dtype=torch.int32, - device='cuda').unsqueeze(0) - torch.testing.assert_close(model_engine.model.recorded_position_ids, - expected_prefill_pos_ids, - atol=0, - rtol=0) - - # Simulate decoding one token after prefill - requests[-1].add_new_token(42, 0) - - # Generation run - batch = ScheduledRequests() - batch.generation_requests = requests - _prepare_kv_cache_resources(kv_cache_manager, batch) - - model_engine.forward(batch, resource_manager) - expected_gen_pos_id = torch.tensor([prompt_len], - dtype=torch.int32, - device='cuda').unsqueeze(0) - torch.testing.assert_close(model_engine.model.recorded_position_ids, - expected_gen_pos_id, - atol=0, - rtol=0) - - kv_cache_manager.shutdown() - - def test_warmup(self) -> None: - model_engine, kv_cache_manager = create_model_engine_and_kvcache(self) - resource_manager = ResourceManager( - {ResourceManagerType.KV_CACHE_MANAGER: kv_cache_manager}) - - # Test with a huge batch size. The warmup run should bail out of - # warmup instead of crashing (there's not enough KV cache space for this). - model_engine._runner._config.cuda_graph_batch_sizes.append(1000000000) - - num_free_before = kv_cache_manager.get_num_free_blocks() - model_engine.warmup(resource_manager) - - # Warmup pre-allocates the CUDA graph padding dummy for each captured - # draft length (only draft_len 0 here, no speculation), which - # intentionally keeps holding its KV blocks so padding cannot fall - # back to eager once the cache saturates. - padding_dummies = model_engine.cuda_graph_runner.padding_dummy_requests - self.assertEqual([0], sorted(padding_dummies.keys())) - - # Make sure we don't leak any blocks beyond that dummy: freeing it - # must restore the exact pre-warmup free-block count. - kv_cache_manager.free_resources(padding_dummies[0]) - self.assertEqual(num_free_before, - kv_cache_manager.get_num_free_blocks()) - - kv_cache_manager.shutdown() - - def test_warmup_skips_padding_dummy_when_padding_impossible(self) -> None: - # With max_batch_size=1 and a graph for batch size 1, every batch - # already matches a graph size, so no padded batch can ever occur and - # warmup must not retain a padding dummy (it would permanently hold - # KV blocks and spec/hybrid resources that are never used). - llm_args = TorchLlmArgs(model="dummy", - max_batch_size=1, - max_num_tokens=258, - cuda_graph_config=CudaGraphConfig( - enable_padding=True, batch_sizes=[1, 2, 4])) - model_engine, kv_cache_manager = create_model_engine_and_kvcache( - self, llm_args) - resource_manager = ResourceManager( - {ResourceManagerType.KV_CACHE_MANAGER: kv_cache_manager}) - - num_free_before = kv_cache_manager.get_num_free_blocks() - model_engine.warmup(resource_manager) - - self.assertEqual({}, - model_engine.cuda_graph_runner.padding_dummy_requests) - self.assertEqual(num_free_before, - kv_cache_manager.get_num_free_blocks()) - - kv_cache_manager.shutdown() - - def test_warmup_skips_padding_dummy_during_estimation(self) -> None: - # The estimation-phase KV cache is sized with no headroom for retained - # dummies; holding blocks there can leave the estimation requests - # unschedulable. Warmup must skip the preallocation for managers - # created for KV cache capacity estimation. - model_engine, kv_cache_manager = create_model_engine_and_kvcache(self) - kv_cache_manager.is_estimating_kv_cache = True - resource_manager = ResourceManager( - {ResourceManagerType.KV_CACHE_MANAGER: kv_cache_manager}) - - num_free_before = kv_cache_manager.get_num_free_blocks() - model_engine.warmup(resource_manager) - - self.assertEqual({}, - model_engine.cuda_graph_runner.padding_dummy_requests) - self.assertEqual(num_free_before, - kv_cache_manager.get_num_free_blocks()) - - kv_cache_manager.shutdown() - - def test_preallocate_padding_dummies_uses_captured_draft_lens(self) -> None: - # Speculative engines pad batches with the runtime draft length, so - # the preallocation must create dummies for the draft lengths of the - # captured graphs — not draft_len 0 unconditionally (an extra - # draft_len-0 dummy would permanently hold KV blocks that runtime - # padding never uses). - model_engine, kv_cache_manager = create_model_engine_and_kvcache(self) - resource_manager = ResourceManager( - {ResourceManagerType.KV_CACHE_MANAGER: kv_cache_manager}) - - runner = model_engine.cuda_graph_runner - # Simulate a spec-decode capture: graphs keyed (batch_size, draft_len, - # short_seq_len_mode, is_all_greedy_sample) exist only for draft_len 2. - runner.graphs[(8, 2, False, True)] = Mock() - try: - runner.preallocate_padding_dummies(resource_manager) - - self.assertEqual([2], sorted(runner.padding_dummy_requests.keys())) - finally: - runner.graphs.clear() - for dummy in runner.padding_dummy_requests.values(): - kv_cache_manager.free_resources(dummy) - kv_cache_manager.shutdown() - - def test_release_padding_dummy_covers_every_manager(self) -> None: - # A padding dummy's request ID is spread across several managers, so - # releasing only the main KV cache manager leaves the others holding - # it — and re-creation reuses the same ID. - model_engine, kv_cache_manager = create_model_engine_and_kvcache(self) - spec_manager = Mock() - cross_manager = Mock() - resource_manager = ResourceManager({ - ResourceManagerType.KV_CACHE_MANAGER: - kv_cache_manager, - ResourceManagerType.SPEC_RESOURCE_MANAGER: - spec_manager, - ResourceManagerType.CROSS_KV_CACHE_MANAGER: - cross_manager, - }) - - runner = model_engine.cuda_graph_runner - try: - self.assertIsNotNone( - runner._get_or_create_padding_dummy(resource_manager, 0)) - dummy = runner.padding_dummy_requests[0] - - self.assertTrue(runner.release_padding_dummy(resource_manager, 0)) - - # Dropped from the runner so the lazy path re-creates it... - self.assertEqual({}, runner.padding_dummy_requests) - # ...and the spec resource manager slot is released too, not just - # the main KV cache manager. - spec_manager.free_resources.assert_called_once_with(dummy) - # The cross-KV manager is only involved for encoder-decoder, which - # this engine is not. - self.assertFalse(runner.is_encoder_decoder) - cross_manager.free_resources.assert_not_called() - - # Releasing again is a no-op rather than a double free. - self.assertFalse(runner.release_padding_dummy(resource_manager, 0)) - spec_manager.free_resources.assert_called_once() - finally: - for dummy in runner.padding_dummy_requests.values(): - kv_cache_manager.free_resources(dummy) - runner.padding_dummy_requests.clear() - kv_cache_manager.shutdown() - - def test_layerwise_nvtx_marker(self) -> None: - llm_args = TorchLlmArgs( - model="dummy", - enable_layerwise_nvtx_marker=True, - cuda_graph_config=CudaGraphConfig(enable_padding=True)) - model_engine, kv_cache_manager = create_model_engine_and_kvcache( - self, llm_args) - resource_manager = ResourceManager( - {ResourceManagerType.KV_CACHE_MANAGER: kv_cache_manager}) - - prompt_len = 32 - requests = [_create_request(prompt_len, 0)] - - batch = ScheduledRequests() - batch.context_requests_last_chunk = requests - _prepare_kv_cache_resources(kv_cache_manager, batch) - model_engine.forward(batch, resource_manager) - - expected_prefill_pos_ids = torch.arange(0, - prompt_len, - dtype=torch.int32, - device='cuda').unsqueeze(0) - torch.testing.assert_close(model_engine.model.recorded_position_ids, - expected_prefill_pos_ids, - atol=0, - rtol=0) - - kv_cache_manager.shutdown() - - def test_cuda_graph_padding_filters_huge_batch_size(self): - llm_args = TorchLlmArgs( - model="dummy", - cuda_graph_config=CudaGraphConfig( - enable_padding=True, - batch_sizes=[1, 2, 3, 1000000000000000000000000])) - model_engine = DummyModelEngine(llm_args, torch.half) - - self.assertEqual(model_engine.cuda_graph_batch_sizes, - [1, 2, 3, model_engine.max_seq_len]) - self.assertTrue(model_engine.cuda_graph_padding_enabled) - - def test_forward_pass_callable_on_cuda_graph_on(self) -> None: - llm_args = TorchLlmArgs(model="dummy", - cuda_graph_config=CudaGraphConfig( - enable_padding=True, )) - model_engine, kv_cache_manager = create_model_engine_and_kvcache( - self, llm_args) - - mock_callable = Mock() - model_engine.register_forward_pass_callable(mock_callable) - - resource_manager = ResourceManager( - {ResourceManagerType.KV_CACHE_MANAGER: kv_cache_manager}) - - prompt_len = 32 - requests = [_create_request(prompt_len, 0)] - - batch = ScheduledRequests() - batch.context_requests_last_chunk = requests - _prepare_kv_cache_resources(kv_cache_manager, batch) - model_engine.forward(batch, resource_manager) - - mock_callable.assert_called_once() - - def test_forward_pass_callable_on_cuda_graph_off(self) -> None: - model_engine, kv_cache_manager = create_model_engine_and_kvcache(self) - - mock_callable = Mock() - model_engine.register_forward_pass_callable(mock_callable) - - resource_manager = ResourceManager( - {ResourceManagerType.KV_CACHE_MANAGER: kv_cache_manager}) - - prompt_len = 32 - requests = [_create_request(prompt_len, 0)] - - batch = ScheduledRequests() - batch.context_requests_last_chunk = requests - _prepare_kv_cache_resources(kv_cache_manager, batch) - model_engine.forward(batch, resource_manager) - - mock_callable.assert_called_once() - - def test_foward_pass_callable_off(self) -> None: - model_engine, kv_cache_manager = create_model_engine_and_kvcache(self) - self.assertTrue(model_engine._runner.forward_pass_callable is None, - "forward_pass_callback should be None by default") - - # Assert we can run `forward` without a forward_pass_callback without error - resource_manager = ResourceManager( - {ResourceManagerType.KV_CACHE_MANAGER: kv_cache_manager}) - - prompt_len = 32 - requests = [_create_request(prompt_len, 0)] - - batch = ScheduledRequests() - batch.context_requests_last_chunk = requests - _prepare_kv_cache_resources(kv_cache_manager, batch) - model_engine.forward(batch, resource_manager) - - def test_foward_pass_callable_backward_compat(self) -> None: - model_engine, kv_cache_manager = create_model_engine_and_kvcache(self) - self.assertTrue(model_engine._runner.forward_pass_callable is None, - "forward_pass_callback should be None by default") - - # Assert we can run `forward` without a forward_pass_callback without error - resource_manager = ResourceManager( - {ResourceManagerType.KV_CACHE_MANAGER: kv_cache_manager}) - - prompt_len = 32 - requests = [_create_request(prompt_len, 0)] - - batch = ScheduledRequests() - batch.context_requests_last_chunk = requests - _prepare_kv_cache_resources(kv_cache_manager, batch) - model_engine.forward(batch, resource_manager) - - @skip_ray - def test_prepare_inputs_with_helix_parallelism(self) -> None: - """Test prepare_inputs with helix parallelism.""" - - # Create model engine with helix parallelism. - llm_args = TorchLlmArgs(model="dummy") - model_engine = DummyModelEngine(llm_args, dtype=torch.half) - - # Provide mapping for model engine. - cp_size = 2 - cp_rank = 0 - cp_config = {"cp_type": CpType.HELIX, "tokens_per_block": 4} - mapping = Mapping(world_size=cp_size, - tp_size=1, - pp_size=1, - cp_size=cp_size, - cp_config=cp_config, - rank=cp_rank) - model_engine._runner._inputs.mapping = mapping - - # Create scheduled requests with two generation requests. - scheduled_requests = ScheduledRequests() - scheduled_requests.context_requests_last_chunk = [] - prompt_lens = [20, 15] - gen_requests = [] - for idx in range(len(prompt_lens)): - req = _create_request(num_tokens=prompt_lens[idx], req_id=idx + 1) - req.py_prompt_len = prompt_lens[idx] - req.py_batch_idx = None - req.is_dummy_request = False - req.py_seq_slot = idx - req.sampling_config.beam_width = 1 - req.py_multimodal_data = {} - req.total_input_len_cp = prompt_lens[idx] * 2 - req.seqlen_this_rank_cp = prompt_lens[idx] - req.py_decoding_iter = 1 - gen_requests.append(req) - scheduled_requests.generation_requests = gen_requests - - # Create KV cache manager for attention metadata. - kv_cache_config = KvCacheConfig(max_tokens=512, - enable_block_reuse=False) - kv_cache_manager = KVCacheManagerV2( - kv_cache_config, - tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF, - num_layers=1, - num_kv_heads=16, - head_dim=16, - tokens_per_block=1, - max_seq_len=512, - max_batch_size=4, - mapping=mapping, - dtype=tensorrt_llm.bindings.DataType.HALF, - ) - kv_cache_manager.add_dummy_requests([1, 2], prompt_lens, is_gen=True) - attn_metadata = TrtllmAttentionMetadata( - max_num_requests=4, - max_num_tokens=512, - kv_cache_manager=kv_cache_manager, - mapping=mapping) - attn_metadata.is_cuda_graph = False - - # Initialize model engine buffers. - max_num_tokens = 512 - model_engine._runner._inputs._config = replace( - model_engine._runner._inputs._config, max_num_tokens=max_num_tokens) - _resize_input_buffers(model_engine._runner, max_num_tokens) - - result = model_engine._runner._inputs.prepare_inputs( - scheduled_requests=scheduled_requests, - kv_cache_manager=kv_cache_manager, - attn_metadata=attn_metadata, - enable_spec_decode=False, - runtime_draft_len=0, - is_dummy=False).kwargs - - # Verify expected keys are present. - self.assertIsNotNone(result) - self.assertIn('input_ids', result) - self.assertIn('position_ids', result) - self.assertIn('attn_metadata', result) - - # Also, verify that position_ids are properly calculated. - position_ids = result['position_ids'] - self.assertIsInstance(position_ids, torch.Tensor) - expected_positions = [40, 30] - actual_positions = position_ids.squeeze(0).cpu().tolist()[:2] - self.assertEqual( - actual_positions, expected_positions, - f"Position IDs should reflect CP allgather results. Expected: {expected_positions}, Got: {actual_positions}" - ) - - # Verify attention metadata is properly configured. - self.assertEqual(attn_metadata.request_ids, [1, 2]) - self.assertEqual(attn_metadata.prompt_lens, [20, 15]) - self.assertEqual(attn_metadata.num_contexts, 0) - - # Verify KV cache parameters - self.assertIsNotNone(attn_metadata.kv_cache_params) - self.assertTrue(attn_metadata.kv_cache_params.use_cache) - - # Verify sequence lengths are correct. - expected_seq_lens = [1, 1] - if hasattr(attn_metadata, - 'seq_lens') and attn_metadata.seq_lens is not None: - actual_seq_lens = attn_metadata.seq_lens.cpu().tolist() - self.assertEqual(actual_seq_lens, expected_seq_lens) - - def _setup_mrope_engine(self, - max_num_tokens: int = 32, - spec_config: DecodingBaseConfig | None = None): - """Build a DummyModelEngine that takes the MRoPE path, plus its KV cache - manager and attention metadata.""" - llm_args = TorchLlmArgs(model="dummy") - model_engine = DummyModelEngine(llm_args, - dtype=torch.half, - spec_config=spec_config) - model_engine.model.model_config.pretrained_config.rope_scaling = { - "type": "mrope" - } - _set_runner_config(model_engine._runner, use_mrope=True) - - mapping = Mapping(world_size=1, tp_size=1, rank=0) - kv_cache_config = KvCacheConfig(max_tokens=max_num_tokens) - kv_cache_manager = KVCacheManagerV2( - kv_cache_config, - tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF, - num_layers=1, - num_kv_heads=16, - head_dim=16, - tokens_per_block=1, - max_seq_len=max_num_tokens, - max_batch_size=4, - mapping=mapping, - dtype=tensorrt_llm.bindings.DataType.HALF, - ) - attn_metadata = AttentionMetadata(max_num_requests=4, - max_num_tokens=max_num_tokens, - kv_cache_manager=kv_cache_manager) - attn_metadata.is_cuda_graph = False - - _set_runner_config(model_engine._runner, max_num_tokens=max_num_tokens) - _resize_input_buffers(model_engine._runner, max_num_tokens) - return model_engine, kv_cache_manager, attn_metadata - - @staticmethod - def _make_mrope_gen_request(num_tokens: int, req_id: int, seq_slot: int, - delta): - """Generation request carrying an MRoPE delta when `delta` is not None, - and no multimodal data at all otherwise (a text-only prompt).""" - request = _create_request(num_tokens, req_id) - request.py_prompt_len = num_tokens - request.py_batch_idx = None - request.py_seq_slot = seq_slot - request.sampling_config.beam_width = 1 - if delta is None: - request.py_multimodal_data = {} - else: - request.py_multimodal_data = { - "mrope_config": { - "mrope_position_deltas": - torch.tensor([[delta]], dtype=torch.int32) - }, - } - return request - - def _make_mrope_spec_batch(self, previous_batch: bool): - """Two multimodal verify groups around a text-only one, plus a CUDA graph - padding dummy, each carrying one draft token. - - Returns the scheduled batch and the multimodal requests.""" - # (num_tokens, req_id, seq_slot, delta); the middle request is text-only. - requests = [ - self._make_mrope_gen_request(4, 1, 0, 10), - self._make_mrope_gen_request(5, 2, 1, None), - self._make_mrope_gen_request(6, 3, 2, 20), - ] - for request in requests: - request.py_draft_tokens = [7] - if previous_batch: - # The overlap scheduler's previous step ran these requests at - # their current slots, which also seeded the delta cache. - request.py_batch_idx = request.py_seq_slot - if request.py_multimodal_data: - request.py_mrope_delta_cache_slot = request.py_seq_slot - # The padding dummy carries the zero delta that `add_dummy_requests` - # attaches on MRoPE models. - dummy = self._make_mrope_gen_request(2, 4, 3, 0) - dummy.is_cuda_graph_dummy = True - dummy.py_draft_tokens = [0] - - scheduled_requests = ScheduledRequests() - scheduled_requests.context_requests_last_chunk = [] - scheduled_requests.generation_requests = requests + [dummy] - return scheduled_requests, [requests[0], requests[2]] - - def test_prepare_inputs_spec_decode_adds_mrope_delta(self) -> None: - """One-model speculative verify groups get the request's MRoPE delta. - - Every token a generation step feeds is text after the prompt's vision - spans, so its three coordinates are `seq_index + delta`, as for an - ordinary generation request. Read slots stay dense, one per verify - group, and a generation-only request seeds the delta cache. - """ - model_engine, kv_cache_manager, attn_metadata = self._setup_mrope_engine( - spec_config=SADecodingConfig(max_draft_len=1)) - scheduled_requests, multimodal_requests = self._make_mrope_spec_batch( - previous_batch=False) - spec_metadata = Mock( - _force_non_greedy_for_capture=False, - context_prompt_lookahead_tokens=None, - ) - - result = model_engine._runner._inputs.prepare_inputs( - scheduled_requests=scheduled_requests, - kv_cache_manager=kv_cache_manager, - attn_metadata=attn_metadata, - spec_metadata=spec_metadata, - enable_spec_decode=True, - runtime_draft_len=1, - is_dummy=False).kwargs - - # Each verify group is [last token, draft token] at past_seen + 0/1; - # the text-only request and the zero-delta dummy keep the scalar - # positions. - position_ids = result["position_ids"] - self.assertEqual(tuple(position_ids.shape), (3, 1, 8)) - expected = torch.tensor([[[13, 14, 4, 5, 25, 26, 1, 2]]] * 3, - dtype=torch.int32, - device='cuda') - torch.testing.assert_close(position_ids, expected, atol=0, rtol=0) - # Slot 32 is the reserved zero slot (max_num_tokens * pp_size). - self.assertEqual(result["mrope_delta_read_seq_slots"].cpu().tolist(), - [0, 32, 2, 32]) - self.assertEqual(result["mrope_delta_write_seq_slots"].cpu().tolist(), - [0, 2]) - self.assertEqual([ - request.py_mrope_delta_cache_slot for request in multimodal_requests - ], [0, 2]) - kv_cache_manager.shutdown() - - def test_prepare_inputs_spec_decode_overlap_keeps_mrope_delta(self) -> None: - """Under the overlap scheduler the accepted-token correction applied in - `preprocess_inputs` lands on top of the delta, on all three axes.""" - model_engine, kv_cache_manager, attn_metadata = self._setup_mrope_engine( - spec_config=SADecodingConfig(max_draft_len=1)) - scheduled_requests, _ = self._make_mrope_spec_batch(previous_batch=True) - spec_metadata = Mock( - _force_non_greedy_for_capture=False, - context_prompt_lookahead_tokens=None, - ) - # The previous step accepted 2, 1 and 2 tokens in slots 0, 1 and 2. - overlap_state = SampleStateTensorsSpec( - new_tokens=torch.zeros((2, 4, 1), dtype=torch.int32, device="cuda"), - new_tokens_lens=torch.tensor([2, 1, 2, 1], - dtype=torch.int32, - device="cuda"), - next_draft_tokens=torch.zeros((4, 1), - dtype=torch.int32, - device="cuda"), - ) - - inputs = model_engine._runner._inputs.prepare_inputs( - scheduled_requests=scheduled_requests, - kv_cache_manager=kv_cache_manager, - attn_metadata=attn_metadata, - spec_metadata=spec_metadata, - new_tensors_device=overlap_state, - enable_spec_decode=True, - runtime_draft_len=1, - is_dummy=False).kwargs - - # The host packs positions from the stale base (past_seen + delta). - stale = torch.tensor([[[13, 14, 4, 5, 25, 26, 1, 2]]] * 3, - dtype=torch.int32, - device='cuda') - torch.testing.assert_close(inputs["position_ids"], - stale, - atol=0, - rtol=0) - self.assertEqual(inputs["mrope_delta_read_seq_slots"].cpu().tolist(), - [0, 32, 2, 32]) - # The delta cache already holds both multimodal requests' deltas. - self.assertNotIn("mrope_delta_write_seq_slots", inputs) - - inputs = model_engine._runner._inputs.preprocess_inputs( - inputs, enable_spec_decode=True, runtime_draft_len=1) - corrected = torch.tensor([[[15, 16, 5, 6, 27, 28, 1, 2]]] * 3, - dtype=torch.int32, - device='cuda') - torch.testing.assert_close(inputs["position_ids"], - corrected, - atol=0, - rtol=0) - kv_cache_manager.shutdown() - - @staticmethod - def _make_mrope_ctx_request(num_tokens: int, req_id: int, seq_slot: int, - delta: int): - """Single-chunk context request carrying MRoPE metadata. - - Its position on axis `a` of token `i` is `100 + 1000 * a + i`, so a - spliced span says which axis and token it came from.""" - request = _create_request(num_tokens, req_id) - request.py_prompt_len = num_tokens - request.py_batch_idx = None - request.py_seq_slot = seq_slot - request.context_current_position = 0 - request.context_chunk_size = num_tokens - request.sampling_config.beam_width = 1 - positions = (100 + 1000 * torch.arange(3).view(3, 1, 1) + - torch.arange(num_tokens).view(1, 1, -1)).to(torch.int32) - request.py_multimodal_data = { - "mrope_config": { - "mrope_position_ids": positions, - "mrope_position_deltas": torch.tensor([[delta]], - dtype=torch.int32), - }, - } - return request - - @staticmethod - def _prepare_mrope_spec_inputs(model_engine, kv_cache_manager, - attn_metadata, scheduled_requests): - spec_metadata = Mock( - _force_non_greedy_for_capture=False, - context_prompt_lookahead_tokens=None, - ) - result = model_engine._runner._inputs.prepare_inputs( - scheduled_requests=scheduled_requests, - kv_cache_manager=kv_cache_manager, - attn_metadata=attn_metadata, - spec_metadata=spec_metadata, - enable_spec_decode=True, - runtime_draft_len=1, - is_dummy=False).kwargs - return result - - def _assert_mrope_positions(self, result, rows) -> None: - """`rows` is one list for all three axes, or one list per axis.""" - if not isinstance(rows[0], list): - rows = [rows] * 3 - position_ids = result["position_ids"] - self.assertEqual(tuple(position_ids.shape), (3, 1, len(rows[0]))) - expected = torch.tensor([rows[0:1], rows[1:2], rows[2:3]], - dtype=torch.int32, - device='cuda') - torch.testing.assert_close(position_ids, expected, atol=0, rtol=0) - - def test_prepare_inputs_spec_decode_mrope_plain_rows_follow_verify_groups( - self) -> None: - """Verify groups are packed ahead of plain generation requests, and the - MRoPE delta read slots follow that packing rather than scheduler order. - - The attention kernel indexes the deltas by generation batch index: row - i of the packed generation batch must read the delta of the request - packed i-th, so a verify group's slot comes before every plain - request's, whatever order the scheduler listed them in. - """ - model_engine, kv_cache_manager, attn_metadata = self._setup_mrope_engine( - spec_config=MTPDecodingConfig(max_draft_len=1)) - # `verify_*` carry a draft token, `plain_*` do not; `*_mm` have an - # MRoPE delta, `*_text` none. - plain_mm = self._make_mrope_gen_request(6, 1, 0, 20) - verify_mm = self._make_mrope_gen_request(4, 2, 1, 10) - plain_text = self._make_mrope_gen_request(5, 3, 2, None) - verify_text = self._make_mrope_gen_request(3, 4, 3, None) - verify_mm.py_draft_tokens = [7] - verify_text.py_draft_tokens = [7] - - scheduled_requests = ScheduledRequests() - scheduled_requests.context_requests_last_chunk = [] - scheduled_requests.generation_requests = [ - plain_mm, verify_mm, plain_text, verify_text - ] - result = self._prepare_mrope_spec_inputs(model_engine, kv_cache_manager, - attn_metadata, - scheduled_requests) - - # Packed order: verify_mm, verify_text, plain_mm, plain_text. The - # verify groups hold two rows each, the plain requests one. - self._assert_mrope_positions(result, [13, 14, 2, 3, 25, 4]) - # Slot 32 is the reserved zero slot (max_num_tokens * pp_size) read by - # the text-only requests. - self.assertEqual(result["mrope_delta_read_seq_slots"].cpu().tolist(), - [1, 32, 0, 32]) - self.assertEqual(result["mrope_delta_write_seq_slots"].cpu().tolist(), - [1, 0]) - kv_cache_manager.shutdown() - - def test_prepare_inputs_spec_decode_mrope_context_takes_no_read_slot( - self) -> None: - """A context request packs ahead of the verify groups with its own - MRoPE positions and takes no generation read slot, so the verify - groups' slots start the list.""" - model_engine, kv_cache_manager, attn_metadata = self._setup_mrope_engine( - spec_config=MTPDecodingConfig(max_draft_len=1)) - # A context request with multimodal content takes the placeholder-token - # split, which reads the vocabulary size. - model_engine.model.config.vocab_size = 1000 - context = self._make_mrope_ctx_request(5, 1, 0, 30) - verify_mm = self._make_mrope_gen_request(4, 2, 1, 10) - verify_mm.py_draft_tokens = [7] - plain_mm = self._make_mrope_gen_request(6, 3, 2, 20) - - scheduled_requests = ScheduledRequests() - scheduled_requests.context_requests_last_chunk = [context] - scheduled_requests.generation_requests = [plain_mm, verify_mm] - result = self._prepare_mrope_spec_inputs(model_engine, kv_cache_manager, - attn_metadata, - scheduled_requests) - - # The context span keeps the processor's per-axis positions; the - # generation rows after it are text at position + delta on all axes. - self._assert_mrope_positions(result, [ - [100, 101, 102, 103, 104, 13, 14, 25], - [1100, 1101, 1102, 1103, 1104, 13, 14, 25], - [2100, 2101, 2102, 2103, 2104, 13, 14, 25], - ]) - self.assertEqual(result["mrope_delta_read_seq_slots"].cpu().tolist(), - [1, 2]) - # The context request seeds the delta cache first, then the verify - # group and the plain request seed theirs. - self.assertEqual(result["mrope_delta_write_seq_slots"].cpu().tolist(), - [0, 1, 2]) - self.assertEqual(context.py_mrope_delta_cache_slot, 0) - kv_cache_manager.shutdown() - - def _setup_extend_as_context_engine(self): - """A two-model draft mode on an attention backend that runs verify groups - as chunked context rows. - - `extend_ctx` is False for every one-model mode and for the TRTLLM - backend below sm_120, so this takes a drafter-based mode and swaps in - the vanilla backend. - """ - model_engine, kv_cache_manager, attn_metadata = self._setup_mrope_engine( - spec_config=NGramDecodingConfig(max_draft_len=1)) - _set_runner_config(model_engine._runner, - attention_backend=VanillaAttention, - disable_overlap_scheduler=True) - spec_config = model_engine._runner._config.spec_config - self.assertTrue(spec_config.spec_dec_mode.extend_ctx(VanillaAttention)) - self.assertTrue(spec_config.is_linear_tree) - return model_engine, kv_cache_manager, attn_metadata - - def test_prepare_inputs_mrope_verify_as_context_takes_no_read_slot( - self) -> None: - """When the kernel runs verify groups as chunked context rows, their - rotary comes from `position_ids` alone and only the plain generation - requests take read slots.""" - model_engine, kv_cache_manager, attn_metadata = ( - self._setup_extend_as_context_engine()) - verify_mm = self._make_mrope_gen_request(4, 1, 0, 10) - verify_text = self._make_mrope_gen_request(3, 2, 1, None) - plain_mm = self._make_mrope_gen_request(6, 3, 2, 20) - plain_text = self._make_mrope_gen_request(5, 4, 3, None) - verify_mm.py_draft_tokens = [7] - verify_text.py_draft_tokens = [7] - - scheduled_requests = ScheduledRequests() - scheduled_requests.context_requests_last_chunk = [] - scheduled_requests.generation_requests = [ - verify_mm, verify_text, plain_mm, plain_text - ] - result = self._prepare_mrope_spec_inputs(model_engine, kv_cache_manager, - attn_metadata, - scheduled_requests) - - self._assert_mrope_positions(result, [13, 14, 2, 3, 25, 4]) - self.assertEqual(result["mrope_delta_read_seq_slots"].cpu().tolist(), - [2, 32]) - # The verify group still seeds the delta cache for the plain steps - # that follow it. - self.assertEqual(result["mrope_delta_write_seq_slots"].cpu().tolist(), - [0, 2]) - kv_cache_manager.shutdown() - - def test_prepare_inputs_mrope_verify_as_context_alone_emits_no_read_slots( - self) -> None: - """A batch whose only delta carriers are verify groups run as context - rows has nothing for the generation-indexed read.""" - model_engine, kv_cache_manager, attn_metadata = ( - self._setup_extend_as_context_engine()) - verify_mm = self._make_mrope_gen_request(4, 1, 0, 10) - verify_mm.py_draft_tokens = [7] - plain_text = self._make_mrope_gen_request(5, 2, 1, None) - - scheduled_requests = ScheduledRequests() - scheduled_requests.context_requests_last_chunk = [] - scheduled_requests.generation_requests = [verify_mm, plain_text] - result = self._prepare_mrope_spec_inputs(model_engine, kv_cache_manager, - attn_metadata, - scheduled_requests) - - self._assert_mrope_positions(result, [13, 14, 4]) - self.assertNotIn("mrope_delta_read_seq_slots", result) - self.assertEqual(result["mrope_delta_write_seq_slots"].cpu().tolist(), - [0]) - kv_cache_manager.shutdown() - - @skip_pre_blackwell_unittest - def test_ced_warmup_passes_dummy_context_requests(self) -> None: - for warmup in (False, True): - with self.subTest(warmup=warmup): - engine, manager, metadata = self._setup_mrope_engine(1024) - manager.has_private_swa_suffix = Mock(return_value=True) - engine.model.model_config.pretrained_config.rope_scaling = None - engine.model.model = SimpleNamespace(ced_kv_precompute=True, - decoder_replay_split=20) - _set_runner_config(engine._runner, - use_mrope=False, - disable_overlap_scheduler=True) - engine._runner.is_warmup = warmup - engine._runner._graphs.cuda_graph_runner = SimpleNamespace( - enabled=False) - scheduled = ScheduledRequests() - scheduled.reset_context_requests( - manager.add_dummy_requests([0], - token_nums=[257], - is_gen=False)) - try: - prepared = engine._runner._inputs.prepare_inputs( - scheduled, - manager, - metadata, - enable_spec_decode=False, - runtime_draft_len=0, - is_dummy=warmup) - inputs = prepared.kwargs - self.assertEqual(inputs['context_requests'], - scheduled.context_requests) - self.assertTrue(inputs['context_requests'][0].is_dummy) - self.assertEqual(metadata.seq_lens.tolist(), [257]) - self.assertEqual(inputs['input_ids'].numel(), 257) - finally: - for req in scheduled.context_requests: - manager.free_resources(req) - manager.shutdown() - - @skip_pre_blackwell_unittest - def test_ced_inputs_preserve_physical_rows_and_token_sources(self) -> None: - from tensorrt_llm._torch.pyexecutor.ced_replay import EncoderReplay - - for overlap, mixed in ((False, True), (True, True), (True, False)): - with self.subTest(overlap=overlap, mixed=mixed): - engine, manager, metadata = self._setup_mrope_engine(5001) - manager.has_private_swa_suffix = Mock(return_value=True) - engine.model.model_config.pretrained_config.rope_scaling = None - engine.model.model = SimpleNamespace(ced_kv_precompute=True, - decoder_replay_split=20) - _set_runner_config(engine._runner, - use_mrope=False, - disable_overlap_scheduler=not overlap) - engine._runner.is_warmup = False - engine._runner._graphs.cuda_graph_runner = SimpleNamespace( - enabled=not mixed) - scheduled = ScheduledRequests() - if mixed: - context = _create_request_with_tokens(list(range(5000)), 7) - context.context_current_position = 4000 - context.context_chunk_size = 1000 - context.py_seq_slot = 0 - context.py_ced_replay = EncoderReplay(7, 123, 4000, 3872) - scheduled.context_requests_last_chunk = [context] - for index, (length, previous_index) in enumerate( - ((20, 2), (30, 0))): - req = _create_request(length, 8 + index) - req.py_seq_slot = index + 1 - req.py_batch_idx = previous_index - scheduled.generation_requests.append(req) - new_tokens = torch.tensor([[[222], [999], [111]]], - dtype=torch.int32, - device='cuda') - try: - prepared = engine._runner._inputs.prepare_inputs( - scheduled, - manager, - metadata, - new_tensors_device=SimpleNamespace( - new_tokens=new_tokens) if overlap else None, - enable_spec_decode=False, - runtime_draft_len=0, - is_dummy=False) - inputs = prepared.kwargs - expected_context = list(range(3872, 5000)) if mixed else [] - self.assertEqual( - inputs['input_ids'].cpu().tolist(), - expected_context + ([111, 222] if overlap else [0, 0])) - self.assertEqual( - inputs['position_ids'].view(-1).cpu().tolist(), - expected_context + ([20, 30] if overlap else [19, 29])) - if mixed: - self.assertEqual(inputs['context_requests'], [context]) - self.assertEqual(metadata.seq_lens.tolist(), - [1128, 1, 1]) - self.assertEqual(context.context_current_position, 4000) - self.assertEqual(context.context_chunk_size, 1000) - self.assertEqual(context.cached_tokens, 4000) - if not overlap: - self.assertEqual( - list(metadata.kv_cache_params. - num_cached_tokens_per_seq), [3872, 19, 29]) - else: - self.assertNotIn('context_requests', inputs) - # GPU sampling is one step ahead of CPU token append. - self.assertEqual([ - r.max_beam_num_tokens - for r in scheduled.generation_requests - ], [20, 30]) - finally: - manager.shutdown() - - def test_item_scheduled_chunk_embedding_is_not_retained(self) -> None: - """The per-forward item-scheduler payload must not replace the - request's own multimodal data, or each chunked prefill would keep its - last chunk's joined embedding rows until prefill completes.""" - model_engine, kv_cache_manager = create_model_engine_and_kvcache(self) - model_engine.model.config.vocab_size = 100 - attn_metadata = AttentionMetadata(max_num_requests=4, - max_num_tokens=32, - kv_cache_manager=kv_cache_manager) - attn_metadata.is_cuda_graph = False - - context = _create_request_with_tokens([11, 22, 33, 44], 1) - context.context_chunk_size = 2 - context.py_seq_slot = 0 - context.py_mm_encoder_state = object() - multimodal_data = { - MULTIMODAL_ENCODER_ITEM_METADATA_KEY: - MultimodalItemMetadata([('image', 0)], output_embedding_lengths=[2]) - } - context.py_multimodal_data = multimodal_data - model_engine._mm_item_scheduler = Mock() - model_engine._runner._inputs._mm_item_scheduler = model_engine._mm_item_scheduler - model_engine._mm_item_scheduler.build_multimodal_data_for_llm.return_value = { - "multimodal_embedding": torch.ones((2, 1), dtype=torch.float16), - "multimodal_embedding_is_chunk": True, - } - scheduled_requests = ScheduledRequests() - scheduled_requests.context_requests_chunking = [context] - - model_engine._runner._inputs.prepare_inputs( - scheduled_requests=scheduled_requests, - kv_cache_manager=kv_cache_manager, - attn_metadata=attn_metadata, - enable_spec_decode=False, - runtime_draft_len=0, - is_dummy=False, - ) - - model_engine._mm_item_scheduler.build_multimodal_data_for_llm.assert_called_once( - ) - self.assertIs(context.py_multimodal_data, multimodal_data) - self.assertNotIn("multimodal_embedding", multimodal_data) - kv_cache_manager.shutdown() - - def test_kv_cache_manager_with_execution_stream(self) -> None: - """Test that KVCacheManagerV2 uses the provided execution_stream. - """ - # Create a dedicated execution stream - execution_stream = torch.cuda.Stream() - - model_engine, kv_cache_manager = create_model_engine_and_kvcache( - self, execution_stream=execution_stream) - - # Verify the KVCacheManagerV2 uses the provided execution stream - self.assertEqual( - kv_cache_manager._stream.cuda_stream, execution_stream.cuda_stream, - "KVCacheManagerV2 should use the provided execution_stream") - - resource_manager = ResourceManager( - {ResourceManagerType.KV_CACHE_MANAGER: kv_cache_manager}) - - prompt_len = 32 - requests = [_create_request(prompt_len, 0)] - - batch = ScheduledRequests() - batch.context_requests_last_chunk = requests - _prepare_kv_cache_resources(kv_cache_manager, batch) - with torch.cuda.stream(execution_stream): - model_engine.forward(batch, resource_manager) - - # Verify the stream is still the same after forward pass - self.assertEqual( - kv_cache_manager._stream.cuda_stream, execution_stream.cuda_stream, - "KVCacheManagerV2 should still use the provided execution_stream after forward" - ) - - kv_cache_manager.shutdown() - - def test_cuda_graph_replay_observes_execution_stream_dependency( - self) -> None: - """A graph replay on the KV manager stream waits for restored KV data.""" - execution_stream = torch.cuda.Stream() - transfer_stream = torch.cuda.Stream() - _, kv_cache_manager = create_model_engine_and_kvcache( - self, execution_stream=execution_stream) - - source = torch.zeros(1, dtype=torch.int32, device="cuda") - observed = torch.zeros_like(source) - graph = torch.cuda.CUDAGraph() - torch.cuda.synchronize() - with torch.cuda.graph(graph, stream=execution_stream): - observed.copy_(source) - - ready = torch.cuda.Event() - with torch.cuda.stream(transfer_stream): - # Model the async host-to-device restore completed by the local - # offload manager. Recording the event after the write is the same - # dependency shape that refreshBlocks/resume installs. - source.fill_(7) - ready.record() - - manager_stream = torch.cuda.ExternalStream( - kv_cache_manager._stream.cuda_stream) - with torch.cuda.stream(manager_stream): - manager_stream.wait_event(ready) - graph.replay() - torch.cuda.synchronize() - - self.assertEqual(observed.item(), 7) - kv_cache_manager.shutdown() - - -@pytest.mark.parametrize("draft_capacity", [None, 0, 8]) -def test_runner_input_buffers_support_repeated_cuda_graph_replay( - draft_capacity): - buffers = InputBuffers.allocate( - max_num_tokens=16, - max_batch_size=2, - max_beam_width=2, - max_seq_len=8, - use_cache_indirection=True, - max_num_draft_tokens=draft_capacity, - ) - assert buffers.cache_indirection.shape == (2, 2, 8) - if draft_capacity is None: - assert buffers.gather_ids_cuda is None - assert buffers.draft_tokens_cuda is None - else: - assert buffers.gather_ids_cuda.numel() == 16 - assert buffers.draft_tokens_cuda.numel() == draft_capacity - input_ptr = buffers.input_ids_cuda.data_ptr() - buffers.input_ids_cuda.fill_(1) - buffers.position_ids_cuda.fill_(2) - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(stream): - for _ in range(3): - output = buffers.input_ids_cuda + buffers.position_ids_cuda - torch.cuda.current_stream().wait_stream(stream) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - output = buffers.input_ids_cuda + buffers.position_ids_cuda - output_ptr = output.data_ptr() - for token in (3, 7): - buffers.input_ids_cuda.fill_(token) - graph.replay() - torch.testing.assert_close(output, torch.full_like(output, token + 2)) - assert output.data_ptr() == output_ptr - assert buffers.input_ids_cuda.data_ptr() == input_ptr - graph.reset() - - -class TestDecodeGraphMetadataPreparation(unittest.TestCase): - - @staticmethod - def _metadata(): - from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import \ - CSA2TrtllmMetadata - metadata = object.__new__(CSA2TrtllmMetadata) - metadata.is_cuda_graph = True - return metadata - - def test_scope_requires_selected_ordinary_decode_graph(self): - key = KeyType(batch_size=1, draft_len=0) - cases = ( - "ordinary", - "ordinary_overlap", - "graph_miss", - "eager_metadata", - "context", - "promoted_context", - "speculative", - "beam", - "encoder_decoder", - "context_parallel", - "draft_key", - ) - for case in cases: - with self.subTest(case=case): - engine, _, _, _, _ = _make_forward_only_engine(key) - enable_spec_decode = False - metadata = self._metadata() - batch = ScheduledRequests() - batch.generation_requests = [_make_request_stub(1)] - selected = key - promoted = frozenset() - engine._runner._config.disable_overlap_scheduler = case != "ordinary_overlap" - if case == "graph_miss": - selected = None - elif case == "eager_metadata": - metadata.is_cuda_graph = False - elif case == "context": - batch.context_requests_last_chunk = [_make_request_stub(2)] - elif case == "promoted_context": - promoted = frozenset({1}) - elif case == "speculative": - enable_spec_decode = True - elif case == "beam": - engine._runner._config.max_beam_width = 2 - elif case == "encoder_decoder": - engine._runner._deferred_decode_outputs_supported = False - elif case == "context_parallel": - engine.mapping.cp_size = 2 - elif case == "draft_key": - selected = key._replace(draft_len=1) - with engine._runner._cuda_graph_metadata_prepare_scope( - metadata, - selected, - batch, - promoted, - enable_spec_decode=enable_spec_decode): - self.assertEqual(metadata._csa2_defer_decode_outputs, case - in ("ordinary", "ordinary_overlap")) - self.assertFalse(metadata._csa2_defer_decode_outputs) - - def test_forward_scopes_proof_to_prepare_for_capture_replay_and_fallback( - self): - key = KeyType(batch_size=1, draft_len=0) - for mode in ("warmup", "capture", "replay", "fallback", - "prepare_error"): - with self.subTest(mode=mode): - selected = None if mode == "fallback" else key - engine, runner, resource_manager, _, outputs = _make_forward_only_engine( - selected) - metadata = self._metadata() - engine._runner.attn_metadata = metadata - engine._runner._metadata.setup_attention.return_value = metadata - runner.maybe_get_cuda_graph.return_value = (metadata, None, - selected) - runner.needs_capture.return_value = mode in ("warmup", - "capture") - runner.is_warmup_only = mode == "warmup" - runner.capture.return_value = outputs - batch = ScheduledRequests() - batch.generation_requests = [_make_request_stub(1)] - prepared = engine._runner._inputs.prepare_inputs.return_value - - def prepare(*args, **kwargs): - self.assertIs(metadata.trtllm_gen_jit_warmup, False) - self.assertEqual(metadata._csa2_defer_decode_outputs, mode - != "fallback") - if mode == "prepare_error": - raise RuntimeError("prepare failed") - return prepared - - engine._runner._inputs.prepare_inputs.side_effect = prepare - with patch( - "tensorrt_llm._torch.pyexecutor.model_engine.torch.cuda.Event", - return_value=Mock()): - if mode == "prepare_error": - with self.assertRaisesRegex(RuntimeError, - "prepare failed"): - engine.forward(batch, resource_manager) - else: - actual_outputs = engine.forward(batch, resource_manager) - self.assertEqual(actual_outputs, { - **outputs, "runtime_draft_len": 0 - }) - self.assertIs(actual_outputs["logits"], - outputs["logits"]) - self.assertFalse(metadata._csa2_defer_decode_outputs) - if mode in ("warmup", "capture"): - runner.capture.assert_called_once() - if mode in ("capture", "replay"): - runner.replay.assert_called_once() - if mode == "fallback": - engine._runner._forward_step.assert_called_once() - - def test_jit_warmup_flag_marks_only_the_warmup_forward(self): - engine, _, resource_manager, _, _ = _make_forward_only_engine(None) - runner = engine._runner - metadata = self._metadata() - runner.attn_metadata = metadata - runner._metadata.setup_attention.return_value = metadata - prepared = runner._inputs.prepare_inputs.return_value - observed = [] - - def prepare(*args, **kwargs): - observed.append(metadata.trtllm_gen_jit_warmup) - return prepared - - runner._inputs.prepare_inputs.side_effect = prepare - batch = ScheduledRequests() - batch.generation_requests = [_make_request_stub(1)] - for jit_warmup in (True, False): - runner._forward_warmup(batch, - resource_manager, - enable_spec_decode=False, - runtime_draft_len=0, - trtllm_gen_jit_warmup=jit_warmup) - self.assertEqual(observed, [True, False]) - self.assertFalse(runner._trtllm_gen_jit_warmup) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_draft_len_schedule.py b/tests/unittest/_torch/speculative/hw_agnostic/test_draft_len_schedule.py index 28590a04a7c2..de42b2836132 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_draft_len_schedule.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_draft_len_schedule.py @@ -1,5 +1,22 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 + +from types import SimpleNamespace + +import pytest +import torch +from utils.llm_data import llm_models_root +from utils.util import similar + +from tensorrt_llm import LLM, SamplingParams +from tensorrt_llm._torch.speculative.utils import ( + get_draft_len_for_batch_size, + resolve_draft_len, + update_draft_len, +) +from tensorrt_llm.llmapi import DraftTargetDecodingConfig, KvCacheConfig, NGramDecodingConfig +from tensorrt_llm.llmapi.llm_args import PARDDecodingConfig + # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -13,15 +30,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import pytest -import torch -from utils.llm_data import llm_models_root -from utils.util import similar - -from tensorrt_llm import LLM, SamplingParams -from tensorrt_llm._torch.speculative.utils import get_draft_len_for_batch_size -from tensorrt_llm.llmapi import DraftTargetDecodingConfig, KvCacheConfig, NGramDecodingConfig - # # ============================================================================ # # Fixture: Force single-worker mode (only for tests that use mocking) @@ -345,3 +353,48 @@ def instrumented_prepare_draft(scheduled_batch, resource_manager): assert actual_len <= drafter_tokens, ( f"Iter {idx}, req {req_idx}: NGram produced {actual_len} > max {drafter_tokens}" ) + + +@pytest.mark.parametrize("config_cls", [DraftTargetDecodingConfig, PARDDecodingConfig]) +def test_draft_length_resolution_resyncs_dynamic_buffers(config_cls): + config = config_cls(max_draft_len=3, speculative_model="dummy", draft_len_schedule={1: 3, 4: 2}) + engine = SimpleNamespace( + spec_config=config, + max_draft_len=3, + max_total_draft_tokens=config.tokens_per_gen_step - 1, + ) + requests = [ + SimpleNamespace( + py_draft_tokens=[7, 8, 9] + [0] * (config.tokens_per_gen_step - 4), + py_needs_onehot_draft_probs=False, + ) + for _ in range(2) + ] + batch = SimpleNamespace(batch_size=2, generation_requests=requests) + + def resolve(**kwargs): + return resolve_draft_len(config, batch, max_draft_len=3, static_draft_len=3, **kwargs) + + # Normal iteration selects K=2. + assert update_draft_len(engine, batch) == 2 + assert all(request.py_draft_tokens[:2] == [7, 8] for request in requests) + # Each explicit warmup shape resynchronizes the buffers and returns its + # length, even when the batch size is unchanged. + for draft_len in (0, 3, 1): + assert resolve(draft_len=draft_len) == draft_len + assert all( + len(request.py_draft_tokens) == config.get_runtime_tokens_per_gen_step(draft_len) - 1 + for request in requests + ) + assert resolve() == 2 + assert resolve(speculation_permanently_disabled=True) == 0 + assert all(request.py_draft_tokens == [] for request in requests) + + +def test_static_draft_length_resolution_preserves_proposals(): + request = SimpleNamespace(py_draft_tokens=[7], py_needs_onehot_draft_probs=False) + batch = SimpleNamespace(batch_size=1, generation_requests=[request]) + + assert resolve_draft_len(None, batch, max_draft_len=3, static_draft_len=5) == 5 + assert request.py_draft_tokens == [7] + assert not request.py_needs_onehot_draft_probs