[https://nvbugs/5850094][fix] Fix MoE cost estimation for auto multi-stream scheduling (#11160)

yizhang-nv · web-flow · commit 48206f3c9576 · 2026-02-02T04:28:58.000-05:00
Signed-off-by: Yi Zhang &lt;187001205+yizhang-nv@users.noreply.github.com&gt;
diff --git a/tensorrt_llm/_torch/compilation/multi_stream/auto_multi_stream.py b/tensorrt_llm/_torch/compilation/multi_stream/auto_multi_stream.py
@@ -39,6 +39,7 @@ def estimate_time(node: Node) -> int:
     moe_ops = {
         torch.ops.trtllm.fp4_block_scale_moe_runner.default,
         torch.ops.trtllm.fused_moe.default,
+        torch.ops.trtllm.moe_custom_op.default,
     }
 
     gemm_ops = {

Original file line number	Diff line number	Diff line change
`@@ -39,6 +39,7 @@ def estimate_time(node: Node) -> int:`
`39`	`39`	`moe_ops = {`
`40`	`40`	`torch.ops.trtllm.fp4_block_scale_moe_runner.default,`
`41`	`41`	`torch.ops.trtllm.fused_moe.default,`
	`42`	`+ torch.ops.trtllm.moe_custom_op.default,`
`42`	`43`	`}`
`43`	`44`
`44`	`45`	`gemm_ops = {`