torch-ext/torch_binding.h

#pragma once

#include <torch/torch.h>

// at::Tensor BatchDecodeWithPagedKVCachePlan_swa_false(
//     at::Tensor float_workspace_buffer, 
//     at::Tensor int_workspace_buffer,
//     at::Tensor page_locked_int_workspace_buffer, 
//     at::Tensor indptr, 
//     int64_t batch_size,
//     int64_t num_qo_heads, 
//     int64_t num_kv_heads, 
//     int64_t page_size, 
//     bool enable_cuda_graph,
//     int64_t window_left, 
//     double logits_soft_cap, 
//     int64_t head_dim_qk, 
//     int64_t head_dim_vo,
//     at::Tensor empty_q_data, 
//     at::Tensor empty_kv_data
// );

// void BatchDecodeWithPagedKVCacheRun_swa_false(
//     at::Tensor float_workspace_buffer,
//     at::Tensor int_workspace_buffer, 
//     at::Tensor plan_info_vec,
//     at::Tensor q, 
//     at::Tensor paged_k_cache,
//     at::Tensor paged_v_cache, 
//     at::Tensor paged_kv_indptr,
//     at::Tensor paged_kv_indices, 
//     at::Tensor paged_kv_last_page_len,                  
//     at::Tensor o, 
//     std::optional<at::Tensor> maybe_lse,
//     int64_t kv_layout_code, 
//     int64_t window_left,
//     bool enable_pdl,
//     std::optional<at::Tensor> maybe_alibi_slopes, 
//     double logits_soft_cap, 
//     double sm_scale, 
//     double rope_rcp_scale, 
//     double rope_rcp_theta
// );

// at::Tensor BatchDecodeWithPagedKVCachePlan_swa_true(
//     at::Tensor float_workspace_buffer, 
//     at::Tensor int_workspace_buffer,
//     at::Tensor page_locked_int_workspace_buffer, 
//     at::Tensor indptr, 
//     int64_t batch_size,
//     int64_t num_qo_heads, 
//     int64_t num_kv_heads, 
//     int64_t page_size, 
//     bool enable_cuda_graph,
//     int64_t window_left, 
//     double logits_soft_cap, 
//     int64_t head_dim_qk, 
//     int64_t head_dim_vo,
//     at::Tensor empty_q_data, 
//     at::Tensor empty_kv_data
// );

// void BatchDecodeWithPagedKVCacheRun_swa_true(
//     at::Tensor float_workspace_buffer,
//     at::Tensor int_workspace_buffer, 
//     at::Tensor plan_info_vec,
//     at::Tensor q, 
//     at::Tensor paged_k_cache,
//     at::Tensor paged_v_cache, 
//     at::Tensor paged_kv_indptr,
//     at::Tensor paged_kv_indices, 
//     at::Tensor paged_kv_last_page_len,                  
//     at::Tensor o, 
//     std::optional<at::Tensor> maybe_lse,
//     int64_t kv_layout_code, 
//     int64_t window_left,
//     bool enable_pdl,
//     std::optional<at::Tensor> maybe_alibi_slopes, 
//     double logits_soft_cap, 
//     double sm_scale, 
//     double rope_rcp_scale, 
//     double rope_rcp_theta
// );

void gelu_and_mul(
    at::Tensor& out, 
    at::Tensor& input, 
    bool enable_pdl
);

void gelu_tanh_and_mul(
    at::Tensor& out, 
    at::Tensor& input, 
    bool enable_pdl
);

void silu_and_mul(
    at::Tensor& out, 
    at::Tensor& input, 
    bool enable_pdl
);