647428d8d6
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
36 lines
1.1 KiB
Python
36 lines
1.1 KiB
Python
import torch
|
|
|
|
from sglang.multimodal_gen.runtime.layers.custom_op import CustomOp
|
|
from sglang.multimodal_gen.runtime.layers.triton_ops import fuse_scale_shift_kernel
|
|
|
|
|
|
class MulAdd(CustomOp):
|
|
"""
|
|
Fuse elementwise mul and add
|
|
Input: a, b, c, OptionalInt[k]
|
|
Output: a * (k + b) + c
|
|
"""
|
|
|
|
def __init__(self, prefix: str = ""):
|
|
super().__init__()
|
|
|
|
def forward_native(
|
|
self, a: torch.Tensor, b: torch.Tensor, c: torch.Tensor, k: int = 0
|
|
) -> torch.Tensor:
|
|
# a.shape: [batch_size, seq_len, inner_dim]
|
|
if b.dim() == 4:
|
|
# b.shape: [batch_size, num_frames, 1, inner_dim]
|
|
num_frames = b.shape[1]
|
|
frame_seqlen = a.shape[1] // num_frames
|
|
return c + (
|
|
a.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * b
|
|
).flatten(1, 2)
|
|
else:
|
|
# b.shape: [batch_size, 1, inner_dim]
|
|
return c + a * (k + b)
|
|
|
|
def forward_cuda(
|
|
self, a: torch.Tensor, b: torch.Tensor, c: torch.Tensor, k: int = 0
|
|
):
|
|
return fuse_scale_shift_kernel(a, b, c, scale_constant=k)
|