diff --git a/CMakeLists.txt b/CMakeLists.txt index ec92727cf..46c495ee4 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -263,6 +263,30 @@ if(WITH_TORCH) find_library(C10_LIB c10 HINTS ${_torch_lib_dirs} REQUIRED) set(TORCH_LIBRARIES ${TORCH_LIB} ${TORCH_CPU_LIB} ${C10_LIB}) + if(WITH_MOORE) + execute_process( + COMMAND ${_TORCH_PYTHON} -c + "import pathlib, torch_musa; print(pathlib.Path(torch_musa.__file__).resolve().parent / 'share' / 'cmake' / 'TorchMusa')" + OUTPUT_VARIABLE TorchMusa_DIR + OUTPUT_STRIP_TRAILING_WHITESPACE + RESULT_VARIABLE _torch_musa_result + ) + if(NOT _torch_musa_result EQUAL 0 OR NOT TorchMusa_DIR) + message(FATAL_ERROR "`WITH_MOORE` and `WITH_TORCH` require `torch_musa`.") + endif() + set(_infini_ops_torch_libraries "${TORCH_LIBRARIES}") + find_package(TorchMusa CONFIG REQUIRED) + set(TORCH_LIBRARIES "${_infini_ops_torch_libraries}") + unset(_infini_ops_torch_libraries) + get_filename_component(_torch_musa_lib_dir + "${TorchMusa_LIBRARY}" DIRECTORY) + list(APPEND TORCH_RUNTIME_DIRS "${_torch_musa_lib_dir}") + # Retain TorchMusa's dispatcher registrations. The active Moore block + # links the MUSA toolkit/runtime selected by InfiniOps. + list(APPEND TORCH_LIBRARIES + -Wl,--no-as-needed ${TorchMusa_LIBRARY} -Wl,--as-needed) + endif() + if(WITH_NVIDIA) find_library(TORCH_CUDA_LIB torch_cuda HINTS ${_torch_lib_dirs} REQUIRED) find_library(C10_CUDA_LIB c10_cuda HINTS ${_torch_lib_dirs} REQUIRED) diff --git a/docs/build.md b/docs/build.md index ff235b2c2..8bb4c29aa 100644 --- a/docs/build.md +++ b/docs/build.md @@ -58,6 +58,14 @@ toolchain. Operator-pruned builds require it only when `flash_attn_with_kvcache` is selected; the default smoke subset does not select this operator. +Builds with both `WITH_MOORE=ON` and `WITH_TORCH=ON` require an installed +TorchMusa package that provides `TorchMusaConfig.cmake`. The Moore +`flash_attn_varlen_func` backend uses TorchMusa's ATen FlashAttention kernel. +Runtime use additionally requires the provider to register +`aten::_flash_attention_forward` for PrivateUse1. That provider currently +supports global attention only; causal attention also requires matching query +and key sequence lengths. + The built wheel installs the InfiniOps Python extension and the InfiniRT shared library needed by the extension. diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 1259c09a1..852fb0c17 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -480,7 +480,7 @@ if(WITH_NVIDIA) list(APPEND _infini_ops_smoke_ops cutlass_scaled_mm moe_sum) endif() -if(WITH_NVIDIA AND WITH_TORCH) +if((WITH_NVIDIA OR WITH_MOORE) AND WITH_TORCH) list(APPEND _infini_ops_smoke_ops flash_attn_varlen_func) endif() @@ -618,7 +618,7 @@ if(WITH_TORCH) "${PROJECT_SOURCE_DIR}/generated/torch/*.cpp" ) - if(NOT WITH_NVIDIA) + if(NOT WITH_NVIDIA AND NOT WITH_MOORE) list(FILTER TORCH_SOURCES EXCLUDE REGEX ".*/flash_attn_varlen_func/flash_attn_varlen_func\\.cc$") endif() @@ -626,6 +626,9 @@ if(WITH_TORCH) list(FILTER TORCH_SOURCES EXCLUDE REGEX ".*/flash_attn_with_kvcache/flash_attn\\.cc$") endif() + if(NOT WITH_MOORE) + list(FILTER TORCH_SOURCES EXCLUDE REGEX ".*/torch/moore/.*") + endif() set(INFINI_OPS_TORCH_UNITY_BATCH_SIZE "8" CACHE STRING "Number of torch sources to include in each generated unity translation unit; set to 1 to disable") @@ -717,7 +720,8 @@ if(WITH_TORCH) list(APPEND _torch_vendor_include_flags "-I${MUSA_ROOT}/include") endif() - # Vendor-specific defines required by forked `torch` headers. + # Select vendor-specific generated implementations and any definitions + # required by forked `torch` headers. set(_torch_extra_flags "") if(WITH_METAX) list(APPEND _torch_extra_flags "-DUSE_MACA=1" "-DWITH_METAX=1") diff --git a/src/torch/moore/ops/flash_attn_varlen_func/flash_attn_varlen_func.cc b/src/torch/moore/ops/flash_attn_varlen_func/flash_attn_varlen_func.cc new file mode 100644 index 000000000..4b777e841 --- /dev/null +++ b/src/torch/moore/ops/flash_attn_varlen_func/flash_attn_varlen_func.cc @@ -0,0 +1,40 @@ +#include "torch/moore/ops/flash_attn_varlen_func/flash_attn_varlen_func.h" + +#include +#include + +#include "torch/tensor_.h" + +namespace infini::ops { + +void Operator::operator()( + const Tensor q, const Tensor k, const Tensor v, const Tensor cu_seqlens_q, + const Tensor cu_seqlens_k, const int64_t max_seqlen_q, + const int64_t max_seqlen_k, const double dropout_p, + const std::optional softmax_scale, const bool causal, + const std::vector window_size, const double softcap, + const std::optional alibi_slopes, const bool deterministic, + const bool return_attn_probs, const std::optional block_table, + Tensor out) const { + TORCH_CHECK(window_size[0] < 0 && window_size[1] < 0, + "TorchMusa FlashAttention does not support local windows"); + + if (causal) { + auto at_cu_seqlens_q = ToAtenTensor( + const_cast(cu_seqlens_q.data()), cu_seqlens_q_shape_, + cu_seqlens_q_strides_, cu_seqlens_q_dtype_, device_index_); + auto at_cu_seqlens_k = ToAtenTensor( + const_cast(cu_seqlens_k.data()), cu_seqlens_k_shape_, + cu_seqlens_k_strides_, cu_seqlens_k_dtype_, device_index_); + TORCH_CHECK(at::equal(at_cu_seqlens_q, at_cu_seqlens_k), + "TorchMusa FlashAttention requires matching query and key " + "sequence lengths for causal attention"); + } + + AtenFlashAttnVarlenFunc::operator()( + q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, + dropout_p, softmax_scale, causal, window_size, softcap, alibi_slopes, + deterministic, return_attn_probs, block_table, out); +} + +} // namespace infini::ops diff --git a/src/torch/moore/ops/flash_attn_varlen_func/flash_attn_varlen_func.h b/src/torch/moore/ops/flash_attn_varlen_func/flash_attn_varlen_func.h new file mode 100644 index 000000000..c263b6a37 --- /dev/null +++ b/src/torch/moore/ops/flash_attn_varlen_func/flash_attn_varlen_func.h @@ -0,0 +1,29 @@ +#ifndef INFINI_OPS_TORCH_MOORE_FLASH_ATTN_VARLEN_FUNC_H_ +#define INFINI_OPS_TORCH_MOORE_FLASH_ATTN_VARLEN_FUNC_H_ + +#include "torch/ops/flash_attn_varlen_func/flash_attn_varlen_func.h" + +namespace infini::ops { + +template <> +class Operator + : public AtenFlashAttnVarlenFunc { + public: + using AtenFlashAttnVarlenFunc::AtenFlashAttnVarlenFunc; + using AtenFlashAttnVarlenFunc::operator(); + + void operator()(const Tensor q, const Tensor k, const Tensor v, + const Tensor cu_seqlens_q, const Tensor cu_seqlens_k, + const int64_t max_seqlen_q, const int64_t max_seqlen_k, + const double dropout_p, + const std::optional softmax_scale, const bool causal, + const std::vector window_size, const double softcap, + const std::optional alibi_slopes, + const bool deterministic, const bool return_attn_probs, + const std::optional block_table, + Tensor out) const override; +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_TORCH_MOORE_FLASH_ATTN_VARLEN_FUNC_H_ diff --git a/src/torch/ops/flash_attn_varlen_func/flash_attn_varlen_func.cc b/src/torch/ops/flash_attn_varlen_func/flash_attn_varlen_func.cc index 2569c28ca..d1bbb4e38 100644 --- a/src/torch/ops/flash_attn_varlen_func/flash_attn_varlen_func.cc +++ b/src/torch/ops/flash_attn_varlen_func/flash_attn_varlen_func.cc @@ -1,9 +1,6 @@ #include "torch/ops/flash_attn_varlen_func/flash_attn_varlen_func.h" #include -#include -#include -#include #include @@ -11,7 +8,8 @@ namespace infini::ops { -void Operator::operator()( +template +void AtenFlashAttnVarlenFunc::operator()( const Tensor q, const Tensor k, const Tensor v, const Tensor cu_seqlens_q, const Tensor cu_seqlens_k, const int64_t max_seqlen_q, const int64_t max_seqlen_k, const double dropout_p, @@ -26,50 +24,45 @@ void Operator::operator()( (void)return_attn_probs; (void)block_table; - const auto device_index = static_cast(device_index_); - const c10::cuda::CUDAGuard device_guard{device_index}; + auto at_q = ToAtenTensor(const_cast(q.data()), q_shape_, + q_strides_, q_dtype_, device_index_); + auto at_k = ToAtenTensor(const_cast(k.data()), k_shape_, + k_strides_, k_dtype_, device_index_); + auto at_v = ToAtenTensor(const_cast(v.data()), v_shape_, + v_strides_, v_dtype_, device_index_); + auto at_cu_seqlens_q = ToAtenTensor( + const_cast(cu_seqlens_q.data()), cu_seqlens_q_shape_, + cu_seqlens_q_strides_, cu_seqlens_q_dtype_, device_index_); + auto at_cu_seqlens_k = ToAtenTensor( + const_cast(cu_seqlens_k.data()), cu_seqlens_k_shape_, + cu_seqlens_k_strides_, cu_seqlens_k_dtype_, device_index_); + auto at_out = ToAtenTensor(out.data(), out_shape_, out_strides_, + out_dtype_, device_index_); - const auto run = [&]() { - auto at_q = ToAtenTensor(const_cast(q.data()), - q_shape_, q_strides_, - q_dtype_, device_index_); - auto at_k = ToAtenTensor(const_cast(k.data()), - k_shape_, k_strides_, - k_dtype_, device_index_); - auto at_v = ToAtenTensor(const_cast(v.data()), - v_shape_, v_strides_, - v_dtype_, device_index_); - auto at_cu_seqlens_q = ToAtenTensor( - const_cast(cu_seqlens_q.data()), cu_seqlens_q_shape_, - cu_seqlens_q_strides_, cu_seqlens_q_dtype_, device_index_); - auto at_cu_seqlens_k = ToAtenTensor( - const_cast(cu_seqlens_k.data()), cu_seqlens_k_shape_, - cu_seqlens_k_strides_, cu_seqlens_k_dtype_, device_index_); - auto at_out = ToAtenTensor( - out.data(), out_shape_, out_strides_, out_dtype_, device_index_); + std::optional window_size_left; + std::optional window_size_right; + if (window_size[0] >= 0) { + window_size_left = window_size[0]; + } + if (window_size[1] >= 0) { + window_size_right = window_size[1]; + } + if (causal && + (window_size_left.has_value() || window_size_right.has_value())) { + window_size_right = 0; + } - const std::optional window_size_left = - window_size[0] < 0 ? std::nullopt - : std::optional{window_size[0]}; - const std::optional window_size_right = - causal ? std::optional{0} - : window_size[1] < 0 ? std::nullopt - : std::optional{window_size[1]}; + auto result = at::_flash_attention_forward( + at_q, at_k, at_v, at_cu_seqlens_q, at_cu_seqlens_k, max_seqlen_q, + max_seqlen_k, dropout_p, causal, false, softmax_scale, window_size_left, + window_size_right, std::nullopt, std::nullopt); - auto result = at::_flash_attention_forward( - at_q, at_k, at_v, at_cu_seqlens_q, at_cu_seqlens_k, max_seqlen_q, - max_seqlen_k, dropout_p, causal, false, softmax_scale, window_size_left, - window_size_right, std::nullopt, std::nullopt); - - // ATen owns the returned tensor. Keep the InfiniOps trailing-output ABI by - // copying it into the caller-provided buffer on the selected CUDA stream. - at_out.copy_(std::get<0>(result)); - }; - - const c10::cuda::CUDAStreamGuard stream_guard{ - c10::cuda::getStreamFromExternal(reinterpret_cast(stream_), - device_index)}; - run(); + // ATen owns the returned tensor. Keep the InfiniOps trailing-output ABI by + // copying it into the caller-provided buffer on the current ATen stream. + at_out.copy_(std::get<0>(result)); } +template class AtenFlashAttnVarlenFunc; +template class AtenFlashAttnVarlenFunc; + } // namespace infini::ops diff --git a/src/torch/ops/flash_attn_varlen_func/flash_attn_varlen_func.h b/src/torch/ops/flash_attn_varlen_func/flash_attn_varlen_func.h index faf07ce46..ed0fb8227 100644 --- a/src/torch/ops/flash_attn_varlen_func/flash_attn_varlen_func.h +++ b/src/torch/ops/flash_attn_varlen_func/flash_attn_varlen_func.h @@ -5,9 +5,8 @@ namespace infini::ops { -template <> -class Operator - : public FlashAttnVarlenFunc { +template +class AtenFlashAttnVarlenFunc : public FlashAttnVarlenFunc { public: using FlashAttnVarlenFunc::FlashAttnVarlenFunc; using FlashAttnVarlenFunc::operator(); @@ -24,6 +23,14 @@ class Operator Tensor out) const override; }; +template <> +class Operator + : public AtenFlashAttnVarlenFunc { + public: + using AtenFlashAttnVarlenFunc::AtenFlashAttnVarlenFunc; + using AtenFlashAttnVarlenFunc::operator(); +}; + } // namespace infini::ops #endif // INFINI_OPS_TORCH_FLASH_ATTN_VARLEN_FUNC_H_ diff --git a/tests/test_flash_attn_varlen_func.py b/tests/test_flash_attn_varlen_func.py index 1febffdcc..6f19d3912 100644 --- a/tests/test_flash_attn_varlen_func.py +++ b/tests/test_flash_attn_varlen_func.py @@ -19,6 +19,8 @@ ( ((3, 5), (4, 5), 4, 4, False, (-1, -1), None), ((5, 2), (3, 6), 4, 2, True, (-1, -1), 0.125), + ((3, 5), (4, 5), 4, 2, False, (-1, -1), None), + ((3, 5), (3, 5), 4, 2, True, (-1, -1), 0.125), ((4, 3), (6, 2), 4, 2, False, (2, 1), None), ((4, 3), (6, 2), 4, 2, True, (2, 1), None), ), @@ -46,8 +48,11 @@ def test_flash_attn_varlen_func( rtol, atol, ): - if device != "cuda": - pytest.skip("FlashAttention FA2 requires the NVIDIA backend") + _require_flash_attention_backend(device) + if device == "musa" and window_size != (-1, -1): + pytest.skip("TorchMusa FlashAttention does not support local windows") + if device == "musa" and causal and q_lens != k_lens: + pytest.skip("TorchMusa causal FlashAttention requires matching Q/K lengths") q = torch.randn((sum(q_lens), num_heads, head_dim), dtype=dtype, device=device) k = torch.randn((sum(k_lens), num_kv_heads, head_dim), dtype=dtype, device=device) @@ -91,9 +96,10 @@ def test_flash_attn_varlen_func( torch.testing.assert_close(out, expected, rtol=rtol, atol=atol) -def test_flash_attn_varlen_func_non_default_stream(device, implementation_index): - if device != "cuda": - pytest.skip("non-default CUDA streams require the NVIDIA backend") +def test_flash_attn_varlen_func_current_stream(device, implementation_index): + if device not in ("cuda", "musa"): + pytest.skip("non-default streams require a GPU ATen backend") + _require_flash_attention_backend(device) dtype = torch.float16 q_lens = (3, 5) @@ -104,30 +110,33 @@ def test_flash_attn_varlen_func_non_default_stream(device, implementation_index) cu_seqlens_q = _cumulative_lengths(q_lens, device) cu_seqlens_k = _cumulative_lengths(k_lens, device) out = torch.empty_like(q) - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.current_stream()) + device_module = getattr(torch, device) + stream = device_module.Stream() + stream.wait_stream(device_module.current_stream()) + stream_handle = getattr(stream, f"{device}_stream") - infini.ops.flash_attn_varlen_func( - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max(q_lens), - max(k_lens), - 0.0, - None, - True, - (-1, -1), - 0.0, - None, - False, - False, - None, - out, - stream=stream.cuda_stream, - implementation_index=implementation_index, - ) + with device_module.stream(stream): + infini.ops.flash_attn_varlen_func( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max(q_lens), + max(k_lens), + 0.0, + None, + False, + (-1, -1), + 0.0, + None, + False, + False, + None, + out, + stream=stream_handle, + implementation_index=implementation_index, + ) stream.synchronize() expected = _reference_varlen_attention( @@ -137,13 +146,64 @@ def test_flash_attn_varlen_func_non_default_stream(device, implementation_index) q_lens, k_lens, None, - True, + False, (-1, -1), ) torch.testing.assert_close(out, expected, rtol=2e-3, atol=2e-3) -def test_flash_attn_varlen_func_default_stream(device, implementation_index): +@pytest.mark.parametrize( + "q_lens, k_lens, causal, window_size, error", + ( + ((3, 5), (4, 5), True, (-1, -1), "matching query and key"), + ((3, 5), (4, 5), False, (2, 1), "does not support local windows"), + ), +) +def test_flash_attn_varlen_func_moore_capability_guard( + q_lens, + k_lens, + causal, + window_size, + error, + device, + implementation_index, +): + if device != "musa": + pytest.skip("TorchMusa capability guard requires the Moore backend") + _require_flash_attention_backend(device) + + q = torch.randn((sum(q_lens), 4, 64), dtype=torch.float16, device=device) + k = torch.randn((sum(k_lens), 2, 64), dtype=torch.float16, device=device) + v = torch.randn_like(k) + cu_seqlens_q = _cumulative_lengths(q_lens, device) + cu_seqlens_k = _cumulative_lengths(k_lens, device) + out = torch.empty_like(q) + + with pytest.raises(RuntimeError, match=error): + infini.ops.flash_attn_varlen_func( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max(q_lens), + max(k_lens), + 0.0, + None, + causal, + window_size, + 0.0, + None, + False, + False, + None, + out, + stream=get_stream(q.device), + implementation_index=implementation_index, + ) + + +def test_flash_attn_varlen_func_implicit_current_stream(device, implementation_index): if device != "cuda": pytest.skip("CUDA stream coverage requires the NVIDIA backend") @@ -169,8 +229,6 @@ def test_flash_attn_varlen_func_default_stream(device, implementation_index): implementation_index=implementation_index, ) - torch.cuda.default_stream().synchronize() - snapshot = out.clone() current_stream.synchronize() expected = _reference_varlen_attention( q, @@ -182,12 +240,11 @@ def test_flash_attn_varlen_func_default_stream(device, implementation_index): False, (-1, -1), ) - torch.testing.assert_close(snapshot, expected, rtol=2e-3, atol=2e-3) + torch.testing.assert_close(out, expected, rtol=2e-3, atol=2e-3) def test_flash_attn_varlen_func_defaults(device, implementation_index): - if device != "cuda": - pytest.skip("FlashAttention FA2 requires the NVIDIA backend") + _require_flash_attention_backend(device) q = torch.randn((5, 4, 64), dtype=torch.float16, device=device) k = torch.randn((5, 4, 64), dtype=torch.float16, device=device) @@ -269,6 +326,15 @@ def _cumulative_lengths(lengths, device): return torch.tensor(values, dtype=torch.int32, device=device) +def _require_flash_attention_backend(device): + if device not in ("cuda", "musa"): + pytest.skip("FlashAttention FA2 requires the NVIDIA or Moore backend") + if device == "musa" and not torch._C._dispatch_has_kernel_for_dispatch_key( + "aten::_flash_attention_forward", "PrivateUse1" + ): + pytest.skip("TorchMusa does not provide aten::_flash_attention_forward") + + def _reference_varlen_attention( q, k,