vllm_omni.diffusion.layers.custom_op ¶

CustomOp ¶

Bases: Module

Base class for custom ops. Dispatches the forward method to the appropriate backend.

dispatch_forward ¶

dispatch_forward() -> Callable

forward ¶

forward(*args, **kwargs) -> Any

forward_cuda ¶

forward_cuda(*args, **kwargs)

forward_hip ¶

forward_hip(*args, **kwargs)

forward_musa ¶

forward_musa(*args, **kwargs)

forward_native ¶

forward_native(*args, **kwargs)

PyTorch-native implementation of the forward method. This method is optional. If implemented, it can be used with compilers such as torch.compile or PyTorch XLA. Also, it can be used for testing purposes.

forward_npu ¶

forward_npu(*args, **kwargs)

forward_xpu ¶

forward_xpu(*args, **kwargs)