Function PYBIND11_MODULE

csrc/ops.cpp:75–99

PyBind11 Module Registration

torch::Tensor lm_head_sm90_forward(torch::Tensor hidden,
                                   torch::Tensor weight,
                                   torch::optional<torch::Tensor> bias);
torch::Tensor lm_head_sm90_forward_fp32(torch::Tensor hidden,
                                        torch::Tensor weight,
                                        torch::optional<torch::Tensor> bias);
#endif

#if defined(__CUDACC__) || defined(KERNEL_ALIGN_WITH_CUDA) || defined(KERNEL_ALIGN_WITH_ROCM)
torch::Tensor fused_logp_forward_out(torch::Tensor logits, torch::Tensor token_ids, torch::Tensor output);
torch::Tensor fused_logp_forward_fp32(torch::Tensor logits, torch::Tensor token_ids);
torch::Tensor fused_logp_forward_indexed_out(torch::Tensor logits, torch::Tensor token_ids, torch::Tensor row_indices, torch::Tensor output);
torch::Tensor fused_logp_forward_indexed_fp32(torch::Tensor logits, torch::Tensor token_ids, torch::Tensor row_indices);
torch::Tensor fused_logp_forward_online_out(torch::Tensor logits, torch::Tensor token_ids, torch::Tensor output);
torch::Tensor fused_logp_forward_online_fp32(torch::Tensor logits, torch::Tensor token_ids);
torch::Tensor fused_logp_forward_online_indexed_out(torch::Tensor logits, torch::Tensor token_ids, torch::Tensor row_indices, torch::Tensor output);
torch::Tensor fused_logp_forward_online_indexed_fp32(torch::Tensor logits, torch::Tensor token_ids, torch::Tensor row_indices);
torch::Tensor deterministic_logp_forward(torch::Tensor logits, torch::Tensor token_ids);
torch::Tensor deterministic_logp_forward_out(torch::Tensor logits, torch::Tensor token_ids, torch::Tensor output);
torch::Tensor deterministic_logp_forward_fp32(torch::Tensor logits, torch::Tensor token_ids);
torch::Tensor deterministic_logp_forward_indexed_out(torch::Tensor logits, torch::Tensor token_ids, torch::Tensor row_indices, torch::Tensor output);
torch::Tensor deterministic_logp_forward_indexed_fp32(torch::Tensor logits, torch::Tensor token_ids, torch::Tensor row_indices);
std::vector<torch::Tensor> deterministic_logp_tile_stats(
    torch::Tensor logits,
    int64_t vocab_start,
    int64_t real_vocab,
    int64_t num_tiles);
torch::Tensor deterministic_logp_backward(
    torch::Tensor logits,
    torch::Tensor lse,