A high-throughput and memory-efficient inference and serving engine for LLMs
[Bugfix] Fix trtllm fused allreduce+rms_norm for transformers backend (#45307)
Signed-off-by: Thomas Parnell <tpa@zurich.ibm.com>
Thomas Parnell committed
Jun 16, 2026 at 10:34 UTC
ebf3a6d70521214d01f23baca7b0b4f92944abab
1 file changed
+7
vllm/compilation/passes/fusion/allreduce_rms_fusion.py
+7
@@ -155,6 +155,13 @@ if flashinfer_comm is not None:
155
scale_factor: torch.Tensor | None = None,
156
weight_bias: float = 0.0,
157
) -> None:
158
+ # handle transformers backend passing outer batch dim.
159
+ if allreduce_in.dim() != 2:
160
+ hidden = allreduce_in.shape[-1]
161
+ allreduce_in = allreduce_in.view(-1, hidden)
162
+ residual = residual.view(-1, hidden)
163
+ if norm_out is not None:
164
+ norm_out = norm_out.view(-1, hidden)
165
num_tokens, hidden_size = allreduce_in.shape
166
element_size = allreduce_in.element_size()
167
current_tensor_size = num_tokens * hidden_size * element_size