From 73ed6204d791fa8ebba76aaa6ece8be23a6a102b Mon Sep 17 00:00:00 2001
From: Shixiaowei02 <39303645+Shixiaowei02@users.noreply.github.com>
Date: Thu, 11 Jun 2026 03:38:44 -0700
Subject: [PATCH] [None][fix] AutoDeploy: set enable_spec_decode on ADEngine
 for disagg

ADEngine subclasses the abstract ModelEngine and does not run
PyTorchModelEngine.__init__, so it never set `enable_spec_decode`. After
#14546 added an unguarded `self.model_engine.enable_spec_decode` read in
`_prepare_disagg_gen_transmission_complete` (the disagg generation handoff
path that ADEngine traverses via #14057 AutoDeploy Basic Disagg Support),
AutoDeploy disaggregated runs crash with:
  AttributeError: 'ADEngine' object has no attribute 'enable_spec_decode'

#14546 and #14057 each passed CI independently but conflict semantically
once both are on main. Set `is_spec_decode`/`enable_spec_decode` in
ADEngine.__init__, mirroring PyTorchModelEngine
(enable_spec_decode == spec_config is not None), so ADEngine satisfies the
ModelEngine attribute contract that shared PyExecutor code relies on.

Signed-off-by: Shixiaowei02 <39303645+Shixiaowei02@users.noreply.github.com>
---
 tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py
index 6cee4311beec..b272740deab4 100644
--- a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py
+++ b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py
@@ -494,6 +494,11 @@ def __init__(
         else:
             self.max_total_draft_tokens = 0
 
+        # ADEngine skips PyTorchModelEngine.__init__, so set the spec-decode
+        # flags that shared PyExecutor code expects on a ModelEngine.
+        self.is_spec_decode = self.spec_config is not None
+        self.enable_spec_decode = self.is_spec_decode
+
         # For compatibility with PyTorchModelEngine utilities
         self.batch_size = cache_seq_interface.info.max_batch_size