misc: Temporarily disable POD from AOT wheels (#956)

We currently don't generate the AOT implementations for POD Attention.
This causes missing symbols when loading AOT shared library.

I'll try to work on #791 later this week so that we can have a unified
code path for AOT and JIT.

Author: abcdabcd987

csrc/flashinfer_ops.cu

@@ -284,7 +284,8 @@ TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) {
   m.def("batch_prefill_with_paged_kv_cache_run", BatchPrefillWithPagedKVCacheRun);
 
   // pod-attention
-  m.def("pod_with_kv_cache_tensor", pod_with_kv_cache_tensor);
+  // Temporarily disabled because we don't generate the implementation yet.
+  // m.def("pod_with_kv_cache_tensor", pod_with_kv_cache_tensor);
 
   // quantization
   // GPU packbits operator

setup.py

@@ -243,7 +243,7 @@ def __init__(self, *args, **kwargs) -> None:
         "csrc/batch_prefill.cu",
         "csrc/single_decode.cu",
         "csrc/single_prefill.cu",
-        "csrc/pod.cu",
+        # "csrc/pod.cu",  # Temporarily disabled
         "csrc/flashinfer_ops.cu",
     ]
     kernel_sm90_sources = [