Skip to content

Commit 835447b

Browse files
committed
[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>
1 parent 84b349f commit 835447b

1 file changed

Lines changed: 5 additions & 0 deletions

File tree

tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -494,6 +494,11 @@ def __init__(
494494
else:
495495
self.max_total_draft_tokens = 0
496496

497+
# ADEngine skips PyTorchModelEngine.__init__, so set the spec-decode
498+
# flags that shared PyExecutor code expects on a ModelEngine.
499+
self.is_spec_decode = self.spec_config is not None
500+
self.enable_spec_decode = self.is_spec_decode
501+
497502
# For compatibility with PyTorchModelEngine utilities
498503
self.batch_size = cache_seq_interface.info.max_batch_size
499504

0 commit comments

Comments
 (0)