biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[Neuron] Add custom_ops for neuron backend (#13246)

Browse Source 
Signed-off-by: Liangfu Chen <liangfc@amazon.com>
Co-authored-by: George Novack <gnovack@amazon.com>
Co-authored-by: Aoyu Zhang <aoyuzhan@amazon.com>
This commit is contained in:
Liangfu Chen
2025-02-25 11:47:49 -08:00
committed by GitHub
parent 340e39e387
commit f75aa72732
9 changed files with 346 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
vllm/model_executor/custom_op.py
View File

@@ -59,6 +59,11 @@ class CustomOp(nn.Module):
# PyTorch-native implementation.
return self.forward_native(*args, **kwargs)
def forward_neuron(self, *args, **kwargs):
# By default, we assume that Neuron ops are compatible with the
# PyTorch-native implementation.
return self.forward_native(*args, **kwargs)
def forward_oot(self, *args, **kwargs):
# By default, we assume that OOT ops are compatible with the
# PyTorch-native implementation.
@@ -88,6 +93,8 @@ class CustomOp(nn.Module):
return self.forward_tpu
elif current_platform.is_xpu():
return self.forward_xpu
elif current_platform.is_neuron():
return self.forward_neuron
elif current_platform.is_out_of_tree():
return self.forward_oot
else: