[Build] Avoid building too many extensions (#1624)

2023-11-23 16:31:19 -08:00
parent de23687d16
commit e0c6f556e8
25 changed files with 206 additions and 272 deletions
--- a/tests/kernels/test_cache.py
+++ b/tests/kernels/test_cache.py
@@ -3,7 +3,7 @@ import random
 import pytest
 import torch

-from vllm import cache_ops
+from vllm._C import cache_ops

 DTYPES = [torch.half, torch.bfloat16, torch.float]
 NUM_TOKENS = [83]  # Arbitrary values for testing