diff --git a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py index b633d6a3b625..2dd8233cea25 100644 --- a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py +++ b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py @@ -393,11 +393,6 @@ def __init__( self.llm_args.print_iter_log = reporting_info.print_log self.llm_args.enable_iter_perf_stats = reporting_info.enable_iter_perf_stats self.llm_args.enable_iter_req_stats = reporting_info.enable_iter_req_stats - self.llm_args.stream_interval = 1 - self.llm_args.attention_dp_config = None - self.llm_args.batch_wait_timeout_ms = 0 - self.llm_args.batch_wait_timeout_iters = 0 - self.llm_args.batch_wait_max_tokens_ratio = 0.0 self.llm_args.max_num_tokens = cache_seq_interface.info.max_num_tokens self.llm_args.max_seq_len = cache_seq_interface.info.max_seq_len self.iter_counter = 0 @@ -407,12 +402,22 @@ def __init__( self.enable_attention_dp = dist_config.enable_attention_dp if dist_config else False if ad_config is not None: + self.llm_args.stream_interval = ad_config.stream_interval + self.llm_args.attention_dp_config = ad_config.attention_dp_config + self.llm_args.batch_wait_timeout_ms = ad_config.batch_wait_timeout_ms + self.llm_args.batch_wait_timeout_iters = ad_config.batch_wait_timeout_iters + self.llm_args.batch_wait_max_tokens_ratio = ad_config.batch_wait_max_tokens_ratio self.max_beam_width = ad_config.max_beam_width self.spec_config = ad_config.speculative_config self._disable_overlap_scheduler = ad_config.disable_overlap_scheduler self.llm_args.max_stats_len = ad_config.max_stats_len self._enable_chunked_prefill = getattr(ad_config, "enable_chunked_prefill", False) else: + self.llm_args.stream_interval = 1 + self.llm_args.attention_dp_config = None + self.llm_args.batch_wait_timeout_ms = 0 + self.llm_args.batch_wait_timeout_iters = 0 + self.llm_args.batch_wait_max_tokens_ratio = 0.0 self.max_beam_width = 1 self.spec_config = None self._disable_overlap_scheduler = False diff --git a/tests/unittest/auto_deploy/singlegpu/shim/test_engine.py b/tests/unittest/auto_deploy/singlegpu/shim/test_engine.py index a4ccba2610e2..657d9e08e0fe 100644 --- a/tests/unittest/auto_deploy/singlegpu/shim/test_engine.py +++ b/tests/unittest/auto_deploy/singlegpu/shim/test_engine.py @@ -21,10 +21,12 @@ from tensorrt_llm import SamplingParams from tensorrt_llm._torch.auto_deploy._compat import KvCacheConfig +from tensorrt_llm._torch.auto_deploy.llm_args import LlmArgs from tensorrt_llm._torch.auto_deploy.shim.ad_executor import ADEngine from tensorrt_llm._torch.auto_deploy.shim.demollm import DemoEngine from tensorrt_llm._torch.auto_deploy.shim.interface import CachedSequenceInterface from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests +from tensorrt_llm.llmapi import AttentionDpConfig class TransformerLikeModelwithFakeCachePool(nn.Module): @@ -107,6 +109,82 @@ def test_engine(engine_cls: Type[ADEngine], tokens_per_block: int): cache_seq_interface.shutdown() +def _make_cache_seq_interface(device, max_seq_len=64, max_batch_size=8): + return CachedSequenceInterface( + max_seq_len=max_seq_len, + max_batch_size=max_batch_size, + max_num_tokens=default_max_num_tokens(max_seq_len, max_batch_size), + device=device, + kv_cache_config=KvCacheConfig(tokens_per_block=max_seq_len), + ) + + +def test_ad_engine_propagates_pyexecutor_scheduling_config(): + """ADEngine must copy PyExecutor scheduling/streaming knobs from ad_config. + + Regression: these were hardcoded stubs (stream_interval=1, + attention_dp_config=None, batch_wait_*=0), silently dropping the yaml settings. + """ + if not torch.cuda.is_available(): + pytest.skip("ADEngine construction builds the model on CUDA") + + device = torch.device("cuda") + max_seq_len, max_batch_size = 64, 8 + cache_seq_interface = _make_cache_seq_interface(device, max_seq_len, max_batch_size) + cache_seq_interface.to(device) + + ad_config = LlmArgs( + model="test-model", + max_batch_size=max_batch_size, + max_seq_len=max_seq_len, + max_input_len=32, + backend="_autodeploy", + cuda_graph_config={"max_batch_size": max_batch_size}, + stream_interval=20, + attention_dp_config=AttentionDpConfig( + enable_balance=True, batching_wait_iters=50, timeout_iters=1 + ), + batch_wait_timeout_ms=5.0, + batch_wait_timeout_iters=3, + batch_wait_max_tokens_ratio=0.5, + ) + + try: + engine = ADEngine(get_inference_model, cache_seq_interface, ad_config=ad_config) + + assert engine.llm_args.stream_interval == 20 + assert engine.llm_args.attention_dp_config is not None + assert engine.llm_args.attention_dp_config.enable_balance is True + assert engine.llm_args.attention_dp_config.batching_wait_iters == 50 + assert engine.llm_args.attention_dp_config.timeout_iters == 1 + assert engine.llm_args.batch_wait_timeout_ms == 5.0 + assert engine.llm_args.batch_wait_timeout_iters == 3 + assert engine.llm_args.batch_wait_max_tokens_ratio == 0.5 + finally: + cache_seq_interface.shutdown() + + +def test_ad_engine_scheduling_config_defaults_without_ad_config(): + """Without ad_config, ADEngine falls back to the previous stub defaults.""" + if not torch.cuda.is_available(): + pytest.skip("ADEngine construction builds the model on CUDA") + + device = torch.device("cuda") + cache_seq_interface = _make_cache_seq_interface(device) + cache_seq_interface.to(device) + + try: + engine = ADEngine(get_inference_model, cache_seq_interface) + + assert engine.llm_args.stream_interval == 1 + assert engine.llm_args.attention_dp_config is None + assert engine.llm_args.batch_wait_timeout_ms == 0 + assert engine.llm_args.batch_wait_timeout_iters == 0 + assert engine.llm_args.batch_wait_max_tokens_ratio == 0.0 + finally: + cache_seq_interface.shutdown() + + @pytest.mark.parametrize("tokens_per_block", [0, 2]) def test_demo_engine_sampling(tokens_per_block: int): """Test sampling logic specific to DemoEngine."""