diff --git a/docs/build.md b/docs/build.md index 53aaf234c..4fc70a031 100644 --- a/docs/build.md +++ b/docs/build.md @@ -51,11 +51,11 @@ python -m pip install .[dev] \ --config-settings=cmake.define.WITH_NVIDIA=ON ``` -Full builds with both `WITH_NVIDIA=ON` and `WITH_TORCH=ON` include -`flash_attn_with_kvcache` and require a compatible FlashAttention 2.7 Python -distribution to be installed in the build environment. The distribution must -provide the `flash_attn_2_cuda` shared library for the active PyTorch and CUDA -toolchain. Operator-pruned builds require it only when +Full builds with both `WITH_NVIDIA=ON` and `WITH_LINKED=ON` include +`flash_attn_with_kvcache` and require a compatible FlashAttention Python +distribution in the build environment. The distribution must provide a +`flash_attn_2_cuda` shared library that exports the required C++ ABI for the +active PyTorch and CUDA toolchain. Operator-pruned builds require it only when `flash_attn_with_kvcache` is selected; the default smoke subset does not select this operator. diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 3c0b5edc4..749cd9f82 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -620,73 +620,6 @@ if(WITH_LINKED) "Resolved ${_linked_source_count} linked operator source(s).") endif() -set(_build_flash_attn_with_kvcache FALSE) -if(WITH_NVIDIA AND WITH_TORCH AND - (NOT INFINI_OPS_OPS OR - "flash_attn_with_kvcache" IN_LIST _infini_ops_op_allowlist)) - execute_process( - COMMAND ${_TORCH_PYTHON} -c - "import importlib.metadata; print(importlib.metadata.version('flash-attn'))" - OUTPUT_VARIABLE FLASH_ATTN_VERSION - OUTPUT_STRIP_TRAILING_WHITESPACE - RESULT_VARIABLE _flash_attn_version_result - ERROR_QUIET - ) - if(NOT _flash_attn_version_result EQUAL 0 OR - NOT FLASH_ATTN_VERSION MATCHES "^2\\.7([.+-]|$)") - message(FATAL_ERROR - "`flash_attn_with_kvcache` requires installed FlashAttention 2.7; " - "found `${FLASH_ATTN_VERSION}`.") - endif() - - execute_process( - COMMAND ${_TORCH_PYTHON} -c - "import importlib.util; s = importlib.util.find_spec('flash_attn_2_cuda'); print(s.origin if s else '')" - OUTPUT_VARIABLE FLASH_ATTN_LIBRARY - OUTPUT_STRIP_TRAILING_WHITESPACE - RESULT_VARIABLE _flash_attn_library_result - ERROR_QUIET - ) - if(NOT _flash_attn_library_result EQUAL 0 OR - NOT EXISTS "${FLASH_ATTN_LIBRARY}") - message(FATAL_ERROR - "FlashAttention ${FLASH_ATTN_VERSION} does not provide the " - "`flash_attn_2_cuda` shared library.") - endif() - - execute_process( - COMMAND ${CMAKE_NM} -D -C "${FLASH_ATTN_LIBRARY}" - OUTPUT_VARIABLE _flash_attn_symbols - RESULT_VARIABLE _flash_attn_nm_result - ERROR_QUIET - ) - string(FIND "${_flash_attn_symbols}" "flash::mha_fwd_kvcache(" - _flash_attn_kvcache_symbol) - if(NOT _flash_attn_nm_result EQUAL 0 OR - _flash_attn_kvcache_symbol EQUAL -1) - message(FATAL_ERROR - "FlashAttention ${FLASH_ATTN_VERSION} does not export the " - "required `flash::mha_fwd_kvcache` ABI.") - endif() - - set(_build_flash_attn_with_kvcache TRUE) - get_filename_component(FLASH_ATTN_RUNTIME_DIR - "${FLASH_ATTN_LIBRARY}" DIRECTORY) - get_filename_component(FLASH_ATTN_LIBRARY_NAME - "${FLASH_ATTN_LIBRARY}" NAME) - target_link_directories(infiniops PRIVATE "${FLASH_ATTN_RUNTIME_DIR}") - # The Python extension has no SONAME, so link by its exact basename. - target_link_libraries(infiniops PRIVATE - "-l:${FLASH_ATTN_LIBRARY_NAME}") - - list(APPEND TORCH_RUNTIME_DIRS - "${FLASH_ATTN_RUNTIME_DIR}" "$ORIGIN/..") - set_property(TARGET infiniops APPEND PROPERTY BUILD_RPATH - "${FLASH_ATTN_RUNTIME_DIR}") - message(STATUS - "Found FlashAttention ${FLASH_ATTN_VERSION}: ${FLASH_ATTN_LIBRARY}") -endif() - set(TORCH_SOURCES "") if(WITH_TORCH) # Auto-generate ATen-backed operator wrappers from `scripts/torch_ops.yaml`. @@ -743,10 +676,6 @@ if(WITH_TORCH) list(FILTER TORCH_SOURCES EXCLUDE REGEX ".*/flash_attn_varlen_func/flash_attn_varlen_func\\.cc$") endif() - if(NOT _build_flash_attn_with_kvcache) - list(FILTER TORCH_SOURCES EXCLUDE REGEX - ".*/flash_attn_with_kvcache/flash_attn\\.cc$") - endif() endif() if(_infini_ops_linked_uses_torch) diff --git a/src/linked/torch/nvidia/c10.h b/src/linked/torch/nvidia/c10.h new file mode 100644 index 000000000..a24833303 --- /dev/null +++ b/src/linked/torch/nvidia/c10.h @@ -0,0 +1,28 @@ +#ifndef INFINI_OPS_LINKED_TORCH_NVIDIA_C10_H_ +#define INFINI_OPS_LINKED_TORCH_NVIDIA_C10_H_ + +#include +#include +#include + +#include "linked/torch/c10.h" + +namespace infini::ops::linked::torch { + +template <> +struct C10 { + static constexpr Device::Type kDeviceType = Device::Type::kNvidia; + + using StreamGuard = c10::cuda::CUDAStreamGuard; + + static c10::cuda::CUDAStream GetStreamFromExternal(void* stream, + int device_index) { + return c10::cuda::getStreamFromExternal( + reinterpret_cast(stream), + static_cast(device_index)); + } +}; + +} // namespace infini::ops::linked::torch + +#endif // INFINI_OPS_LINKED_TORCH_NVIDIA_C10_H_ diff --git a/src/linked/torch/nvidia/flash_attn.yaml b/src/linked/torch/nvidia/flash_attn.yaml new file mode 100644 index 000000000..c54d82dbd --- /dev/null +++ b/src/linked/torch/nvidia/flash_attn.yaml @@ -0,0 +1,2 @@ +python_distribution_package: flash-attn +library_glob: flash_attn_2_cuda*.so diff --git a/src/linked/torch/nvidia/ops/flash_attn_with_kvcache/flash_attn.cc b/src/linked/torch/nvidia/ops/flash_attn_with_kvcache/flash_attn.cc new file mode 100644 index 000000000..26cf2c910 --- /dev/null +++ b/src/linked/torch/nvidia/ops/flash_attn_with_kvcache/flash_attn.cc @@ -0,0 +1,50 @@ +#include "linked/torch/nvidia/ops/flash_attn_with_kvcache/flash_attn.h" + +namespace flash { + +std::vector mha_fwd_kvcache( + at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, + std::optional& k, std::optional& v, + std::optional& cache_seqlens, + std::optional& rotary_cos, + std::optional& rotary_sin, + std::optional& cache_batch_idx, + std::optional& cache_leftpad, + std::optional& block_table, + std::optional& alibi_slopes, std::optional& out, + float softmax_scale, bool causal, int window_size_left, + int window_size_right, float softcap, bool rotary_interleaved, + int num_splits); + +} // namespace flash + +namespace infini::ops::linked::torch::nvidia { + +std::vector FlashAttnKvcache::Call( + at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, + std::optional& k, std::optional& v, + std::optional& cache_seqlens, + std::optional& rotary_cos, + std::optional& rotary_sin, + std::optional& cache_batch_idx, + std::optional& cache_leftpad, + std::optional& block_table, + std::optional& alibi_slopes, std::optional& out, + float softmax_scale, bool causal, int window_size_left, + int window_size_right, float softcap, bool rotary_interleaved, + int num_splits) { + return flash::mha_fwd_kvcache( + q, k_cache, v_cache, k, v, cache_seqlens, rotary_cos, rotary_sin, + cache_batch_idx, cache_leftpad, block_table, alibi_slopes, out, + softmax_scale, causal, window_size_left, window_size_right, softcap, + rotary_interleaved, num_splits); +} + +} // namespace infini::ops::linked::torch::nvidia + +namespace infini::ops::linked::torch { + +template class TorchFlashAttnWithKvcache< + ::infini::ops::linked::torch::nvidia::FlashAttnKvcache>; + +} // namespace infini::ops::linked::torch diff --git a/src/linked/torch/nvidia/ops/flash_attn_with_kvcache/flash_attn.h b/src/linked/torch/nvidia/ops/flash_attn_with_kvcache/flash_attn.h new file mode 100644 index 000000000..fddb2f73e --- /dev/null +++ b/src/linked/torch/nvidia/ops/flash_attn_with_kvcache/flash_attn.h @@ -0,0 +1,50 @@ +#ifndef INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_FLASH_ATTN_WITH_KVCACHE_FLASH_ATTN_H_ +#define INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_FLASH_ATTN_WITH_KVCACHE_FLASH_ATTN_H_ + +#include "linked/torch/nvidia/c10.h" +#include "linked/torch/ops/flash_attn_with_kvcache.h" + +namespace infini::ops::linked::torch::nvidia { + +struct FlashAttnKvcache : C10 { + static std::vector Call( + at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, + std::optional& k, std::optional& v, + std::optional& cache_seqlens, + std::optional& rotary_cos, + std::optional& rotary_sin, + std::optional& cache_batch_idx, + std::optional& cache_leftpad, + std::optional& block_table, + std::optional& alibi_slopes, std::optional& out, + float softmax_scale, bool causal, int window_size_left, + int window_size_right, float softcap, bool rotary_interleaved, + int num_splits); +}; + +} // namespace infini::ops::linked::torch::nvidia + +namespace infini::ops::linked::torch { + +extern template class TorchFlashAttnWithKvcache< + ::infini::ops::linked::torch::nvidia::FlashAttnKvcache>; + +} // namespace infini::ops::linked::torch + +namespace infini::ops { + +template <> +class Operator + : public linked::torch::TorchFlashAttnWithKvcache< + linked::torch::nvidia::FlashAttnKvcache> { + public: + using linked::torch::TorchFlashAttnWithKvcache< + linked::torch::nvidia::FlashAttnKvcache>::TorchFlashAttnWithKvcache; + + using linked::torch::TorchFlashAttnWithKvcache< + linked::torch::nvidia::FlashAttnKvcache>::operator(); +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_FLASH_ATTN_WITH_KVCACHE_FLASH_ATTN_H_ diff --git a/src/linked/torch/nvidia/ops/flash_attn_with_kvcache/flash_attn.yaml b/src/linked/torch/nvidia/ops/flash_attn_with_kvcache/flash_attn.yaml new file mode 100644 index 000000000..f15198a47 --- /dev/null +++ b/src/linked/torch/nvidia/ops/flash_attn_with_kvcache/flash_attn.yaml @@ -0,0 +1,4 @@ +library: flash_attn +required_symbols: + - >- + flash::mha_fwd_kvcache(at::Tensor&, at::Tensor const&, at::Tensor const&, std::optional&, std::optional&, std::optional&, std::optional&, std::optional&, std::optional&, std::optional&, std::optional&, std::optional&, std::optional&, float, bool, int, int, float, bool, int) diff --git a/src/linked/torch/ops/flash_attn_with_kvcache.h b/src/linked/torch/ops/flash_attn_with_kvcache.h new file mode 100644 index 000000000..2154f6f8c --- /dev/null +++ b/src/linked/torch/ops/flash_attn_with_kvcache.h @@ -0,0 +1,179 @@ +#ifndef INFINI_OPS_LINKED_TORCH_OPS_FLASH_ATTN_WITH_KVCACHE_H_ +#define INFINI_OPS_LINKED_TORCH_OPS_FLASH_ATTN_WITH_KVCACHE_H_ + +#include +#include +#include +#include + +#include "base/flash_attn_with_kvcache.h" +#include "torch/tensor_.h" + +namespace infini::ops::linked::torch { + +template +class TorchFlashAttnWithKvcache : public ::infini::ops::FlashAttnWithKvcache { + public: + using ::infini::ops::FlashAttnWithKvcache::FlashAttnWithKvcache; + + using ::infini::ops::FlashAttnWithKvcache::operator(); + + void operator()(const Tensor q, Tensor k_cache, Tensor v_cache, + const std::optional k, const std::optional v, + const std::optional rotary_cos, + const std::optional rotary_sin, + const std::optional cache_seqlens, + const std::optional cache_batch_idx, + const std::optional cache_leftpad, + const std::optional block_table, + const std::optional alibi_slopes, + const std::optional softmax_scale, const bool causal, + const std::vector window_size, const double softcap, + const bool rotary_interleaved, const int64_t num_splits, + const bool return_softmax_lse, Tensor out, + std::optional softmax_lse) const override { + Run(q, k_cache, v_cache, k, v, rotary_cos, rotary_sin, cache_seqlens, + std::nullopt, cache_batch_idx, cache_leftpad, block_table, + softmax_scale, causal, window_size, softcap, rotary_interleaved, + alibi_slopes, num_splits, out, softmax_lse); + (void)return_softmax_lse; + } + + void operator()(const Tensor q, Tensor k_cache, Tensor v_cache, + const std::optional k, const std::optional v, + const std::optional rotary_cos, + const std::optional rotary_sin, + const int64_t cache_seqlens, + const std::optional cache_batch_idx, + const std::optional cache_leftpad, + const std::optional block_table, + const std::optional alibi_slopes, + const std::optional softmax_scale, const bool causal, + const std::vector window_size, const double softcap, + const bool rotary_interleaved, const int64_t num_splits, + const bool return_softmax_lse, Tensor out, + std::optional softmax_lse) const override { + Run(q, k_cache, v_cache, k, v, rotary_cos, rotary_sin, std::nullopt, + cache_seqlens, cache_batch_idx, cache_leftpad, block_table, + softmax_scale, causal, window_size, softcap, rotary_interleaved, + alibi_slopes, num_splits, out, softmax_lse); + (void)return_softmax_lse; + } + + private: + void Run(const Tensor q, Tensor k_cache, Tensor v_cache, + const std::optional k, const std::optional v, + const std::optional rotary_cos, + const std::optional rotary_sin, + const std::optional cache_seqlens, + const std::optional scalar_cache_seqlens, + const std::optional cache_batch_idx, + const std::optional cache_leftpad, + const std::optional block_table, + const std::optional softmax_scale, const bool causal, + const std::vector window_size, const double softcap, + const bool rotary_interleaved, + const std::optional alibi_slopes, const int64_t num_splits, + Tensor out, std::optional softmax_lse) const { + const typename Backend::StreamGuard stream_guard{ + Backend::GetStreamFromExternal(stream_, device_index_)}; + + auto at_q = ToAtenTensor(const_cast(q.data()), + q_shape_, q_strides_, + q_dtype_, device_index_); + auto at_k_cache = ToAtenTensor( + k_cache.data(), k_cache_shape_, k_cache_strides_, k_cache_dtype_, + device_index_); + auto at_v_cache = ToAtenTensor( + v_cache.data(), v_cache_shape_, v_cache_strides_, v_cache_dtype_, + device_index_); + auto at_out = ToAtenTensor( + out.data(), out_shape_, out_strides_, out_dtype_, device_index_); + std::optional at_softmax_lse; + if (softmax_lse.has_value()) { + at_softmax_lse.emplace(ToAtenTensor( + softmax_lse->data(), softmax_lse_shape_, softmax_lse_strides_, + softmax_lse_dtype_, device_index_)); + } + + std::optional at_k; + std::optional at_v; + std::optional at_rotary_cos; + std::optional at_rotary_sin; + std::optional at_cache_seqlens; + std::optional at_cache_batch_idx; + std::optional at_cache_leftpad; + std::optional at_block_table; + std::optional at_alibi_slopes; + + if (k.has_value()) { + at_k.emplace(ToAtenTensor( + const_cast(k->data()), k_shape_, k_strides_, k_dtype_, + device_index_)); + at_v.emplace(ToAtenTensor( + const_cast(v->data()), v_shape_, v_strides_, v_dtype_, + device_index_)); + } + if (rotary_cos.has_value()) { + at_rotary_cos.emplace(ToAtenTensor( + const_cast(rotary_cos->data()), rotary_cos_shape_, + rotary_cos_strides_, rotary_cos_dtype_, device_index_)); + at_rotary_sin.emplace(ToAtenTensor( + const_cast(rotary_sin->data()), rotary_sin_shape_, + rotary_sin_strides_, rotary_sin_dtype_, device_index_)); + } + if (cache_seqlens.has_value()) { + at_cache_seqlens.emplace(ToAtenTensor( + const_cast(cache_seqlens->data()), cache_seqlens_shape_, + cache_seqlens_strides_, cache_seqlens_dtype_, device_index_)); + } else if (scalar_cache_seqlens.has_value()) { + at_cache_seqlens.emplace( + at::full({static_cast(batch_size_)}, *scalar_cache_seqlens, + at::TensorOptions().dtype(at::kInt).device(at_q.device()))); + } + if (cache_batch_idx.has_value()) { + at_cache_batch_idx.emplace(ToAtenTensor( + const_cast(cache_batch_idx->data()), cache_batch_idx_shape_, + cache_batch_idx_strides_, cache_batch_idx_dtype_, device_index_)); + } + if (cache_leftpad.has_value()) { + at_cache_leftpad.emplace(ToAtenTensor( + const_cast(cache_leftpad->data()), cache_leftpad_shape_, + cache_leftpad_strides_, cache_leftpad_dtype_, device_index_)); + } + if (block_table.has_value()) { + at_block_table = ToAtenTensor( + const_cast(block_table->data()), block_table_shape_, + block_table_strides_, block_table_dtype_, device_index_); + } + if (alibi_slopes.has_value()) { + at_alibi_slopes = ToAtenTensor( + const_cast(alibi_slopes->data()), alibi_slopes_shape_, + alibi_slopes_strides_, alibi_slopes_dtype_, device_index_); + } + + std::optional at_out_optional; + auto result = Backend::Call( + at_q, at_k_cache, at_v_cache, at_k, at_v, at_cache_seqlens, + at_rotary_cos, at_rotary_sin, at_cache_batch_idx, at_cache_leftpad, + at_block_table, at_alibi_slopes, at_out_optional, + static_cast(softmax_scale.value_or( + 1.0 / std::sqrt(static_cast(head_size_)))), + causal, static_cast(window_size[0]), + static_cast(window_size[1]), static_cast(softcap), + rotary_interleaved, static_cast(num_splits)); + assert(!result.empty() && + "Linked `flash_attn_with_kvcache` provider returned no output."); + at_out.copy_(result[0]); + if (at_softmax_lse.has_value()) { + assert(result.size() >= 2 && + "Linked `flash_attn_with_kvcache` provider did not return " + "softmax LSE."); + at_softmax_lse->copy_(result[1]); + } + } +}; + +} // namespace infini::ops::linked::torch + +#endif // INFINI_OPS_LINKED_TORCH_OPS_FLASH_ATTN_WITH_KVCACHE_H_ diff --git a/src/torch/ops/flash_attn_with_kvcache/flash_attn.cc b/src/torch/ops/flash_attn_with_kvcache/flash_attn.cc deleted file mode 100644 index 51f077d3b..000000000 --- a/src/torch/ops/flash_attn_with_kvcache/flash_attn.cc +++ /dev/null @@ -1,191 +0,0 @@ -#include "torch/ops/flash_attn_with_kvcache/flash_attn.h" - -#include -#include -#include - -#include -#include -#include -#include - -#include "torch/tensor_.h" - -namespace flash { - -std::vector mha_fwd_kvcache( - at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, - std::optional& k, std::optional& v, - std::optional& cache_seqlens, - std::optional& rotary_cos, - std::optional& rotary_sin, - std::optional& cache_batch_idx, - std::optional& cache_leftpad, - std::optional& block_table, - std::optional& alibi_slopes, std::optional& out, - float softmax_scale, bool causal, int window_size_left, - int window_size_right, float softcap, bool rotary_interleaved, - int num_splits); - -} // namespace flash - -namespace infini::ops { - -void Operator::operator()( - const Tensor q, Tensor k_cache, Tensor v_cache, - const std::optional k, const std::optional v, - const std::optional rotary_cos, - const std::optional rotary_sin, - const std::optional cache_seqlens, - const std::optional cache_batch_idx, - const std::optional cache_leftpad, - const std::optional block_table, - const std::optional alibi_slopes, - const std::optional softmax_scale, const bool causal, - const std::vector window_size, const double softcap, - const bool rotary_interleaved, const int64_t num_splits, - const bool return_softmax_lse, Tensor out, - std::optional softmax_lse) const { - Run(q, k_cache, v_cache, k, v, rotary_cos, rotary_sin, cache_seqlens, - std::nullopt, cache_batch_idx, cache_leftpad, block_table, softmax_scale, - causal, window_size, softcap, rotary_interleaved, alibi_slopes, - num_splits, out, softmax_lse); - (void)return_softmax_lse; -} - -void Operator::operator()( - const Tensor q, Tensor k_cache, Tensor v_cache, - const std::optional k, const std::optional v, - const std::optional rotary_cos, - const std::optional rotary_sin, const int64_t cache_seqlens, - const std::optional cache_batch_idx, - const std::optional cache_leftpad, - const std::optional block_table, - const std::optional alibi_slopes, - const std::optional softmax_scale, const bool causal, - const std::vector window_size, const double softcap, - const bool rotary_interleaved, const int64_t num_splits, - const bool return_softmax_lse, Tensor out, - std::optional softmax_lse) const { - Run(q, k_cache, v_cache, k, v, rotary_cos, rotary_sin, std::nullopt, - cache_seqlens, cache_batch_idx, cache_leftpad, block_table, softmax_scale, - causal, window_size, softcap, rotary_interleaved, alibi_slopes, - num_splits, out, softmax_lse); - (void)return_softmax_lse; -} - -void Operator::Run( - const Tensor q, Tensor k_cache, Tensor v_cache, - const std::optional k, const std::optional v, - const std::optional rotary_cos, - const std::optional rotary_sin, - const std::optional cache_seqlens, - const std::optional scalar_cache_seqlens, - const std::optional cache_batch_idx, - const std::optional cache_leftpad, - const std::optional block_table, - const std::optional softmax_scale, const bool causal, - const std::vector window_size, const double softcap, - const bool rotary_interleaved, const std::optional alibi_slopes, - const int64_t num_splits, Tensor out, - std::optional softmax_lse) const { - const auto device_index = static_cast(device_index_); - const c10::cuda::CUDAGuard device_guard{device_index}; - const c10::cuda::CUDAStreamGuard stream_guard{ - c10::cuda::getStreamFromExternal(reinterpret_cast(stream_), - device_index)}; - - auto at_q = - ToAtenTensor(const_cast(q.data()), q_shape_, - q_strides_, q_dtype_, device_index_); - auto at_k_cache = ToAtenTensor( - k_cache.data(), k_cache_shape_, k_cache_strides_, k_cache_dtype_, - device_index_); - auto at_v_cache = ToAtenTensor( - v_cache.data(), v_cache_shape_, v_cache_strides_, v_cache_dtype_, - device_index_); - auto at_out = ToAtenTensor( - out.data(), out_shape_, out_strides_, out_dtype_, device_index_); - std::optional at_softmax_lse; - if (softmax_lse.has_value()) { - at_softmax_lse.emplace(ToAtenTensor( - softmax_lse->data(), softmax_lse_shape_, softmax_lse_strides_, - softmax_lse_dtype_, device_index_)); - } - - std::optional at_k; - std::optional at_v; - std::optional at_rotary_cos; - std::optional at_rotary_sin; - std::optional at_cache_seqlens; - std::optional at_cache_batch_idx; - std::optional at_cache_leftpad; - std::optional at_block_table; - std::optional at_alibi_slopes; - - if (k.has_value()) { - at_k.emplace(ToAtenTensor( - const_cast(k->data()), k_shape_, k_strides_, k_dtype_, - device_index_)); - at_v.emplace(ToAtenTensor( - const_cast(v->data()), v_shape_, v_strides_, v_dtype_, - device_index_)); - } - if (rotary_cos.has_value()) { - at_rotary_cos.emplace(ToAtenTensor( - const_cast(rotary_cos->data()), rotary_cos_shape_, - rotary_cos_strides_, rotary_cos_dtype_, device_index_)); - at_rotary_sin.emplace(ToAtenTensor( - const_cast(rotary_sin->data()), rotary_sin_shape_, - rotary_sin_strides_, rotary_sin_dtype_, device_index_)); - } - if (cache_seqlens.has_value()) { - at_cache_seqlens.emplace(ToAtenTensor( - const_cast(cache_seqlens->data()), cache_seqlens_shape_, - cache_seqlens_strides_, cache_seqlens_dtype_, device_index_)); - } else if (scalar_cache_seqlens.has_value()) { - at_cache_seqlens.emplace( - at::full({static_cast(batch_size_)}, *scalar_cache_seqlens, - at::TensorOptions().dtype(at::kInt).device(at_q.device()))); - } - if (cache_batch_idx.has_value()) { - at_cache_batch_idx.emplace(ToAtenTensor( - const_cast(cache_batch_idx->data()), cache_batch_idx_shape_, - cache_batch_idx_strides_, cache_batch_idx_dtype_, device_index_)); - } - if (cache_leftpad.has_value()) { - at_cache_leftpad.emplace(ToAtenTensor( - const_cast(cache_leftpad->data()), cache_leftpad_shape_, - cache_leftpad_strides_, cache_leftpad_dtype_, device_index_)); - } - if (block_table.has_value()) { - at_block_table = ToAtenTensor( - const_cast(block_table->data()), block_table_shape_, - block_table_strides_, block_table_dtype_, device_index_); - } - if (alibi_slopes.has_value()) { - at_alibi_slopes = ToAtenTensor( - const_cast(alibi_slopes->data()), alibi_slopes_shape_, - alibi_slopes_strides_, alibi_slopes_dtype_, device_index_); - } - - std::optional at_out_optional; - auto result = flash::mha_fwd_kvcache( - at_q, at_k_cache, at_v_cache, at_k, at_v, at_cache_seqlens, at_rotary_cos, - at_rotary_sin, at_cache_batch_idx, at_cache_leftpad, at_block_table, - at_alibi_slopes, at_out_optional, - static_cast(softmax_scale.value_or( - 1.0 / std::sqrt(static_cast(head_size_)))), - causal, static_cast(window_size[0]), - static_cast(window_size[1]), static_cast(softcap), - rotary_interleaved, static_cast(num_splits)); - assert(!result.empty() && "`flash::mha_fwd_kvcache` returned no output"); - at_out.copy_(result[0]); - if (at_softmax_lse.has_value()) { - assert(result.size() >= 2 && - "`flash::mha_fwd_kvcache` did not return softmax LSE"); - at_softmax_lse->copy_(result[1]); - } -} - -} // namespace infini::ops diff --git a/src/torch/ops/flash_attn_with_kvcache/flash_attn.h b/src/torch/ops/flash_attn_with_kvcache/flash_attn.h deleted file mode 100644 index 8799bbb81..000000000 --- a/src/torch/ops/flash_attn_with_kvcache/flash_attn.h +++ /dev/null @@ -1,64 +0,0 @@ -#ifndef INFINI_OPS_TORCH_FLASH_ATTN_WITH_KVCACHE_FLASH_ATTN_H_ -#define INFINI_OPS_TORCH_FLASH_ATTN_WITH_KVCACHE_FLASH_ATTN_H_ - -#include "base/flash_attn_with_kvcache.h" - -namespace infini::ops { - -template <> -class Operator - : public FlashAttnWithKvcache { - public: - using FlashAttnWithKvcache::FlashAttnWithKvcache; - using FlashAttnWithKvcache::operator(); - - void operator()(const Tensor q, Tensor k_cache, Tensor v_cache, - const std::optional k, const std::optional v, - const std::optional rotary_cos, - const std::optional rotary_sin, - const std::optional cache_seqlens, - const std::optional cache_batch_idx, - const std::optional cache_leftpad, - const std::optional block_table, - const std::optional alibi_slopes, - const std::optional softmax_scale, const bool causal, - const std::vector window_size, const double softcap, - const bool rotary_interleaved, const int64_t num_splits, - const bool return_softmax_lse, Tensor out, - std::optional softmax_lse) const override; - - void operator()(const Tensor q, Tensor k_cache, Tensor v_cache, - const std::optional k, const std::optional v, - const std::optional rotary_cos, - const std::optional rotary_sin, - const int64_t cache_seqlens, - const std::optional cache_batch_idx, - const std::optional cache_leftpad, - const std::optional block_table, - const std::optional alibi_slopes, - const std::optional softmax_scale, const bool causal, - const std::vector window_size, const double softcap, - const bool rotary_interleaved, const int64_t num_splits, - const bool return_softmax_lse, Tensor out, - std::optional softmax_lse) const override; - - private: - void Run(const Tensor q, Tensor k_cache, Tensor v_cache, - const std::optional k, const std::optional v, - const std::optional rotary_cos, - const std::optional rotary_sin, - const std::optional cache_seqlens, - const std::optional scalar_cache_seqlens, - const std::optional cache_batch_idx, - const std::optional cache_leftpad, - const std::optional block_table, - const std::optional softmax_scale, const bool causal, - const std::vector window_size, const double softcap, - const bool rotary_interleaved, - const std::optional alibi_slopes, const int64_t num_splits, - Tensor out, std::optional softmax_lse) const; -}; - -} // namespace infini::ops - -#endif // INFINI_OPS_TORCH_FLASH_ATTN_WITH_KVCACHE_FLASH_ATTN_H_