fix flashifner enable disable calculation

xinli-sw · xinli-sw · commit 3fb686fa6801 · 2025-07-21T13:50:01.000-07:00
Signed-off-by: XIn Li &lt;xinli@nvidia.com&gt;
diff --git a/vllm/compilation/collective_fusion.py b/vllm/compilation/collective_fusion.py
@@ -176,7 +176,7 @@ def call_trtllm_fused_allreduce_norm(
         use_flashinfer = allreduce_in.shape[0] * allreduce_in.shape[
             1] * allreduce_in.element_size() <= min(
                 _FI_MAX_SIZES[world_size],
-                max_token_num * allreduce_in.shape[0] *
+                max_token_num * allreduce_in.shape[1] *
                 allreduce_in.element_size(),
             )
         if use_flashinfer:

Original file line number	Diff line number	Diff line change
`@@ -176,7 +176,7 @@ def call_trtllm_fused_allreduce_norm(`
`176`	`176`	`use_flashinfer = allreduce_in.shape[0] * allreduce_in.shape[`
`177`	`177`	`1] * allreduce_in.element_size() <= min(`
`178`	`178`	`_FI_MAX_SIZES[world_size],`
`179`		`- max_token_num * allreduce_in.shape[0] *`
	`179`	`+ max_token_num * allreduce_in.shape[1] *`
`180`	`180`	`allreduce_in.element_size(),`
`181`	`181`	`)`
`182`	`182`	`if use_flashinfer:`