diff --git a/python/sglang/multimodal_gen/runtime/layers/elementwise.py b/python/sglang/multimodal_gen/runtime/layers/elementwise.py index 4c5e50857..16db42b90 100644 --- a/python/sglang/multimodal_gen/runtime/layers/elementwise.py +++ b/python/sglang/multimodal_gen/runtime/layers/elementwise.py @@ -23,7 +23,7 @@ class MulAdd(CustomOp): num_frames = b.shape[1] frame_seqlen = a.shape[1] // num_frames return c + ( - a.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * b + a.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * (k + b) ).flatten(1, 2) else: # b.shape: [batch_size, 1, inner_dim]