[yechank-nvidia]

Description

PyExecutor.__init__ wires AsyncTransferManager into the KV cache V2 scheduler's deadlock detector by feature detection:

if hasattr(self.scheduler, "set_async_transfer_manager"):
    self.scheduler.set_async_transfer_manager(self.async_transfer_manager)

When MM encoder item scheduling is enabled, _util.create_py_executor_instance wraps the LLM scheduler in MultimodalScheduler / MultimodalEagerEncoderScheduler. The wrapper did not expose set_async_transfer_manager, so the hasattr check was false and the wrapped KVCacheV2Scheduler._async_transfer_manager stayed None. On a disaggregated context server running KVCacheManagerV2 with MM item scheduling, finished-but-unreaped KV sends (which still pin KV pages) could then be counted as a stall and _detect_deadlock could raise a false deadlock.

Fix: add an explicit set_async_transfer_manager passthrough on MultimodalScheduler (inherited by MultimodalEagerEncoderScheduler) that forwards to the wrapped scheduler when it supports the hook; it is a no-op for wrapped schedulers without it (e.g. SimpleScheduler).

set_async_transfer_manager is the only optional method PyExecutor feature-detects on self.scheduler. The others it calls (schedule_request, can_schedule, is_request_in_schedulable_state, scheduling_state_range) are on the RequestScheduler base and the wrapper already implements or delegates them.

Test Coverage

Recorded validation for de080e435d covers wrapper forwarding; no full disaggregated serving reproduction is claimed.

• tests/unittest/_torch/executor/test_multimodal_scheduler.py::test_multimodal_scheduler_forwards_async_transfer_manager (new, parametrized over MultimodalScheduler and MultimodalEagerEncoderScheduler): checks that the wrapped KVCacheV2Scheduler receives the transfer manager, and that wrapping a scheduler without the hook doesn't fail. It fails without the fix.
• These existing tests still pass: tests/unittest/_torch/executor/test_multimodal_scheduler.py, test_py_scheduler.py and test_scheduler_serializable_output.py.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Configuration used: Repository: NVIDIA/TensorRT-LLM/.coderabbit.yaml
• Review profile: CHILL
• Plan: Enterprise
• Run ID: 406f1c9b-ddf2-43de-98a1-f42a0d1cf3c4

📥 Commits

Reviewing files that changed from the base of the PR and between db359c2 and de080e4.

📒 Files selected for processing (2)

• tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py
• tests/unittest/_torch/executor/test_multimodal_scheduler.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.

Walkthrough

MultimodalScheduler adds a setter that forwards an async transfer manager to the wrapped scheduler when the wrapped scheduler provides the setter. Tests cover both multimodal scheduler classes with wrapped schedulers that provide or lack the hook.

Changes

Multimodal scheduler forwarding

Layer / File(s)	Summary
Async transfer manager setter and tests tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py, tests/unittest/_torch/executor/test_multimodal_scheduler.py	MultimodalScheduler forwards the manager when the wrapped scheduler exposes set_async_transfer_manager. Parameterized tests cover both multimodal scheduler classes and wrapped schedulers with and without the hook.

Priority: ➖ Normal

Estimated code review effort: 2 (Simple) | ~8 minutes

Change: Bug fix

Suggested reviewers: bowenfu

Merge Risk: ⚪ Minimal · up to de080

The multimodal scheduler forwards the transfer manager to the V2 scheduler, while schedulers without the hook remain supported. No merge-blocking risk remains.

Pre-merge checks: 4 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 20.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 5 functions across 2 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks: 4 passed

Check name	Status	Explanation
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.
Description check		The description clearly explains the issue, the passthrough fix, affected schedulers, and relevant test coverage. It is mostly complete despite not reproducing or explicitly checking every checklist i…
Title check		The title clearly and concisely describes the main change: forwarding the async transfer manager through the multimodal scheduler wrapper.

Full details: Description check

Explanation

The description clearly explains the issue, the passthrough fix, affected schedulers, and relevant test coverage. It is mostly complete despite not reproducing or explicitly checking every checklist item.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

_{Comment @coderabbitai help to get the list of available commands.}