[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/benchmarks/kernels/benchmark_paged_attention.py
+++ b/benchmarks/kernels/benchmark_paged_attention.py
@@ -5,7 +5,7 @@ from typing import Optional

 import torch

-from vllm._C import ops
+from vllm import _custom_ops as ops
 from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE, create_kv_caches_with_random

 NUM_BLOCKS = 1024