[Misc] Add indirection layer for custom ops (#3913)

2024-04-11 03:26:07 +00:00
parent e42df7227d
commit e9da5a40c6
14 changed files with 224 additions and 32 deletions
--- a/vllm/model_executor/layers/layernorm.py
+++ b/vllm/model_executor/layers/layernorm.py
@@ -4,7 +4,7 @@ from typing import Optional, Tuple, Union
 import torch
 import torch.nn as nn

-from vllm._C import ops
+from vllm import _custom_ops as ops


 class RMSNorm(nn.Module):