torch-ext/flashinfer/__init__.py

# NOTE: Torch needs to be imported before the custom
# extensions. Otherwise libc10.so cannot be found.
import torch
from ._ops import ops

def gelu_and_mul(out, input, enable_pdl):
    return ops.gelu_and_mul(out, input, enable_pdl) 

def gelu_tanh_and_mul(out, input, enable_pdl):
    return ops.gelu_tanh_and_mul(out, input, enable_pdl)

def silu_and_mul(out, input, enable_pdl):
    return ops.silu_and_mul(out, input, enable_pdl)

# def batch_decode_with_paged_kv_cache_plan_swa_false(
#     float_workspace_buffer,
#     int_workspace_buffer,
#     page_locked_int_workspace_buffer,
#     indptr,
#     batch_size,
#     num_qo_heads,
#     num_kv_heads,
#     page_size,
#     enable_cuda_graph,
#     window_left,
#     logits_soft_cap,
#     head_dim_qk,
#     head_dim_vo,
#     empty_q_data,
#     empty_kv_data
# ):
#     return ops.BatchDecodeWithPagedKVCachePlan_swa_false(
#         float_workspace_buffer,
#         int_workspace_buffer,
#         page_locked_int_workspace_buffer,
#         indptr,
#         batch_size,
#         num_qo_heads,
#         num_kv_heads,
#         page_size,
#         enable_cuda_graph,
#         window_left,
#         logits_soft_cap,
#         head_dim_qk,
#         head_dim_vo,
#         empty_q_data,
#         empty_kv_data
#     )