forked from Karylab-cklius/vllm
46 lines
2.0 KiB
C++
46 lines
2.0 KiB
C++
/* SPDX-License-Identifier: Apache-2.0
|
|
* SPDX-FileCopyrightText: Copyright 2025 SGLang Team. All Rights Reserved.
|
|
*
|
|
* Vendored from sgl-kernel/include/sgl_flash_kernel_ops.h (commit bcf72ccc).
|
|
* Declares the mha_fwd() C++ function signature for CUTLASS FA3 kernels.
|
|
* NO MODIFICATIONS from the original (except removing unused macros).
|
|
*/
|
|
|
|
#pragma once
|
|
|
|
#include <ATen/ATen.h>
|
|
#include <ATen/Tensor.h>
|
|
#include <torch/library.h>
|
|
#include <torch/torch.h>
|
|
|
|
#include <vector>
|
|
|
|
#include "sgl_kernel_torch_shim.h"
|
|
|
|
/*
|
|
* From flash-attention (sgl-attn fork)
|
|
*/
|
|
std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor> mha_fwd(
|
|
at::Tensor q, // (b, s_q, h, d) or (total_q, h, d) if there is cu_seqlens_q
|
|
at::Tensor k, // (b_k, s_k, h_k, d) or (total_k, h_k, d) or paged
|
|
at::Tensor v, // (b_k, s_k, h_k, dv) or (total_k, h_k, dv) or paged
|
|
std::optional<at::Tensor> k_new_, std::optional<at::Tensor> v_new_,
|
|
std::optional<at::Tensor> q_v_, // MLA value projection query
|
|
std::optional<at::Tensor> out_, std::optional<at::Tensor> cu_seqlens_q_,
|
|
std::optional<at::Tensor> cu_seqlens_k_,
|
|
std::optional<at::Tensor> cu_seqlens_k_new_,
|
|
std::optional<at::Tensor> seqused_q_, std::optional<at::Tensor> seqused_k_,
|
|
std::optional<int64_t> max_seqlen_q_, std::optional<int64_t> max_seqlen_k_,
|
|
std::optional<at::Tensor> page_table_,
|
|
std::optional<at::Tensor> kv_batch_idx_,
|
|
std::optional<at::Tensor> leftpad_k_, std::optional<at::Tensor> rotary_cos_,
|
|
std::optional<at::Tensor> rotary_sin_,
|
|
std::optional<at::Tensor> seqlens_rotary_,
|
|
std::optional<at::Tensor> q_descale_, std::optional<at::Tensor> k_descale_,
|
|
std::optional<at::Tensor> v_descale_, std::optional<double> softmax_scale_,
|
|
bool is_causal, int64_t window_size_left, int64_t window_size_right,
|
|
int64_t attention_chunk, double softcap, bool is_rotary_interleaved,
|
|
std::optional<at::Tensor> scheduler_metadata_, int64_t num_splits,
|
|
std::optional<bool> pack_gqa_, int64_t sm_margin,
|
|
std::optional<const at::Tensor>& sinks_);
|