[fix] Fix llama4 min latency (#5117)

liji-nv · web-flow · commit 8282d6c1a70f · 2025-06-11T15:44:38.000+08:00
Signed-off-by: Jin Li &lt;59594262+liji-nv@users.noreply.github.com&gt;
diff --git a/tensorrt_llm/_torch/models/modeling_llama.py b/tensorrt_llm/_torch/models/modeling_llama.py
@@ -301,7 +301,7 @@ def compute_routed_output(self, hidden_states, all_rank_num_tokens,
         routed_output = self.experts(
             hidden_states,
             router_logits,
-            cutlass_min_latency_mode=cutlass_min_latency_mode,
+            do_finalize=not cutlass_min_latency_mode,
             all_rank_num_tokens=all_rank_num_tokens,
             use_dp_padding=use_dp_padding,
         )
diff --git a/tensorrt_llm/_torch/models/modeling_llama_min_latency.py b/tensorrt_llm/_torch/models/modeling_llama_min_latency.py
@@ -515,7 +515,7 @@ def forward(
 
         return super().forward(x,
                                router_logits,
-                               cutlass_min_latency_mode=False,
+                               do_finalize=True,
                                output_dtype=output_dtype)
 
 

Original file line number	Diff line number	Diff line change
`@@ -301,7 +301,7 @@ def compute_routed_output(self, hidden_states, all_rank_num_tokens,`
`301`	`301`	`routed_output = self.experts(`
`302`	`302`	`hidden_states,`
`303`	`303`	`router_logits,`
`304`		`- cutlass_min_latency_mode=cutlass_min_latency_mode,`
	`304`	`+ do_finalize=not cutlass_min_latency_mode,`
`305`	`305`	`all_rank_num_tokens=all_rank_num_tokens,`
`306`	`306`	`use_dp_padding=use_dp_padding,`
`307`	`307`	`)`