[Misc] Add CustomOp interface for device portability

[WoosukKwon]

Currently, the custom layers have two issues. First, they directly import _custom_ops, which are not supported for devices such as TPU and Gaudi. Second, they assume that the custom ops are implemented in the same way for all devices. To address these issues, the PR adds CustomOp interface, an indirection layer that implements the device-specific forward methods. This allows the custom kernels to be lazily imported only for the associated device.

class CustomOp(nn.Module):

    def forward(self, *args, **kwargs):
        if not hasattr(self, "_forward_method"):
            self._forward_method = self.dispatch_forward()
        return self._forward_method(*args, **kwargs)

    def forward_native(self, *args, **kwargs):
        """PyTorch-native implementation of the forward method."""
        raise NotImplementedError

    def forward_cuda(self, *args, **kwargs):
        raise NotImplementedError

    def forward_hip(self, *args, **kwargs):
        ...

    def forward_xpu(self, *args, **kwargs):
        ...

    def forward_cpu(self, *args, **kwargs):
        ...

    def forward_tpu(self, *args, **kwargs):
        ...

    def forward_gaudi(self, *args, **kwargs):
        ...

    def dispatch_forward(self):
        if is_hip():
            return self.forward_hip
        elif is_cpu():
            return self.forward_cpu
        else:
            return self.forward_cuda

According to the benchmarks, the lazy import does not affect the performance:

$ python benchmarks/benchmark_latency.py --model JackFram/llama-68m
# main
Avg latency: 0.3685314357979223 seconds
# This PR
Avg latency: 0.3665724984719418 seconds

[comaniac]

I feel we need more flexibility here in the future. For example, we may build a wheel with both CPU and CUDA enabled, but we want to configure which one to use on the fly. On the other hand, this may not be necessary at this moment.

[WoosukKwon]

Good point. For now, vLLM is bound to a specific backend at the build time. I added a note that we do not support dynamic dispatching currently.

[youkaichao]

pytorch has quite a lot dispatching utilities, can we reuse some? forward usually has a special meaning in pytorch, and it has a special handling in torch.compile . Having these new forward_xxx method might break future torch.compile integration.

[WoosukKwon]

pytorch has quite a lot dispatching utilities, can we reuse some? forward usually has a special meaning in pytorch, and it has a special handling in torch.compile . Having these new forward_xxx method might break future torch.compile integration.

@youkaichao Good point. I moved the dispatching logic to __init__ so that it is not included in the scope of torch.compile. After this change, I believe the PR itself does not break torch.compile.

[comaniac]

LGTM!

[bnellnm]

Hi @WoosukKwon , would you mind taking a look at #5047 before you land this? I've been working on registering all the custom operations via TORCH_LIBRARY (which also has per device dispatching). I'm worried these changes might be at odds with TORCH_LIBRARY/pytorch dispatching.

[WoosukKwon]

@bnellnm Thanks for bringing it up. If I understand correctly, this PR is orthogonal to yours. Basically, I believe your PR does NOT include per-device dispatching, because vLLM always builds the custom library for at most one device. Also, in our situation, dispatching can't be implemented at the C++ level, because we'd like to use Python libraries to implement some custom ops.

[tlrmchlsmth]

lgtm :)