diff --git a/CMakeLists.txt b/CMakeLists.txt index 8152df92..a02f2e52 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -767,8 +767,8 @@ if(ENABLE_NVFP4) endif() # ── Flash-Attention 2 vendored kernels (fp16 + bf16 fwd SM80) ── -# Built as an object library and linked into a SEPARATE pybind module -# flash_rt_fa2.so — same isolation pattern as flash_rt_fp4.so. The +# Built as an object library and linked into a Python-free raw library; +# flash_rt_fa2.so is a thin pybind adapter over that library. The # main flash_rt_kernels.so stays small (~3.6 MB) and its rebuild no # longer waits on FA2's heavy CUTLASS 3.x template codegen (~8 min). # @@ -831,9 +831,9 @@ if(ENABLE_FA2) csrc/attention/fa2_causal_inst/flash_fwd_split_hdim256_bf16_sm80_causal.cu ) endif() - if("bf16" IN_LIST FA2_DTYPES AND ("128" IN_LIST FA2_HDIMS OR "256" IN_LIST FA2_HDIMS)) - list(APPEND FA2_SRCS csrc/attention/fa2_wrapper_causal.cu) - endif() + # Always emit the causal C entry. Builds without a causal hdim retain a + # stable raw-library symbol that fails clearly instead of an unresolved ABI. + list(APPEND FA2_SRCS csrc/attention/fa2_wrapper_causal.cu) add_library(fa2_vendor_obj OBJECT ${FA2_SRCS}) set_target_properties(fa2_vendor_obj PROPERTIES @@ -954,6 +954,22 @@ if(ENABLE_FA2) math(EXPR _FA2_NKERN "${_FA2_NSRC} - 1") message(STATUS "FA2 vendor instantiations: hdim={${_FA2_H}} x dtype={${_FA2_D}} x {split,no-split} = ${_FA2_NKERN} kernel .cu files") + # Keep the raw C entries in a Python-free library. The native C++ runtime + # links this target directly; the pybind module below is only an adapter. + add_library(flashrt_fa2_raw SHARED) + target_sources(flashrt_fa2_raw PRIVATE $) + target_link_libraries(flashrt_fa2_raw PRIVATE CUDA::cudart) + if(NOT MSVC) + target_link_options(flashrt_fa2_raw PRIVATE + "LINKER:--no-undefined" -static-libstdc++ -static-libgcc) + endif() + set_target_properties(flashrt_fa2_raw PROPERTIES + LIBRARY_OUTPUT_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/flash_rt + RUNTIME_OUTPUT_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/flash_rt + BUILD_RPATH "$ORIGIN" + INSTALL_RPATH "$ORIGIN" + BUILD_WITH_INSTALL_RPATH ON) + # Independent pybind module. Caller side: # from flash_rt import flash_rt_fa2 as fa2 # fa2.fwd_fp16(...) / fa2.fwd_bf16(...) @@ -961,12 +977,16 @@ if(ENABLE_FA2) set_target_properties(flash_rt_fa2 PROPERTIES LIBRARY_OUTPUT_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/flash_rt RUNTIME_OUTPUT_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/flash_rt) - target_sources(flash_rt_fa2 PRIVATE $) target_include_directories(flash_rt_fa2 PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/csrc ) - target_link_libraries(flash_rt_fa2 PRIVATE CUDA::cudart) - install(TARGETS flash_rt_fa2 DESTINATION ${CMAKE_CURRENT_SOURCE_DIR}/flash_rt) + target_link_libraries(flash_rt_fa2 PRIVATE flashrt_fa2_raw CUDA::cudart) + set_target_properties(flash_rt_fa2 PROPERTIES + BUILD_RPATH "$ORIGIN" + INSTALL_RPATH "$ORIGIN" + BUILD_WITH_INSTALL_RPATH ON) + install(TARGETS flash_rt_fa2 flashrt_fa2_raw + DESTINATION ${CMAKE_CURRENT_SOURCE_DIR}/flash_rt) message(STATUS "FA2 pybind module: flash_rt_fa2 (separate .so)") endif() @@ -1440,8 +1460,8 @@ if(FLASHRT_ENABLE_MOTUS AND GPU_ARCH STREQUAL "120") ENABLE_TINYFP8_KERNELS=1) endif() -# (ENABLE_FA2 object-lib is linked into flash_rt_fa2.so only — see -# the dedicated pybind11_add_module(flash_rt_fa2 ...) above. The main +# (ENABLE_FA2 object-lib is linked into libflashrt_fa2_raw.so only; the +# dedicated pybind module and native C++ runtime both consume it. The main # flash_rt_kernels.so deliberately does NOT pull FA2 in, so rebuilds # of our hand-written kernels don't re-trigger the FA2 codegen tax.) diff --git a/csrc/attention/fa2_wrapper.cu b/csrc/attention/fa2_wrapper.cu index dd043327..cc8bfcaa 100644 --- a/csrc/attention/fa2_wrapper.cu +++ b/csrc/attention/fa2_wrapper.cu @@ -24,6 +24,7 @@ #include "flash_attn_2_src/flash_attn/namespace_config.h" #include "flash_attn_2_src/flash_attn/flash.h" +#include "attention/fa2_wrapper.h" namespace FLASH_NAMESPACE { template diff --git a/csrc/attention/fa2_wrapper.h b/csrc/attention/fa2_wrapper.h new file mode 100644 index 00000000..d07de0b0 --- /dev/null +++ b/csrc/attention/fa2_wrapper.h @@ -0,0 +1,73 @@ +#ifndef FLASHRT_ATTENTION_FA2_WRAPPER_H +#define FLASHRT_ATTENTION_FA2_WRAPPER_H + +#include + +#ifdef __cplusplus +extern "C" { +#endif + +void fvk_attention_fa2_fwd_fp16( + const void* q_ptr, const void* k_ptr, const void* v_ptr, + void* o_ptr, void* softmax_lse_ptr, + void* softmax_lse_accum_ptr, void* o_accum_ptr, + int batch, int seqlen_q, int seqlen_k, + int num_heads_q, int num_heads_kv, int head_dim, + int q_batch_stride, int q_row_stride, int q_head_stride, + int k_batch_stride, int k_row_stride, int k_head_stride, + int v_batch_stride, int v_row_stride, int v_head_stride, + int o_batch_stride, int o_row_stride, int o_head_stride, + float softmax_scale, int num_sms, cudaStream_t stream); + +void fvk_attention_fa2_fwd_bf16( + const void* q_ptr, const void* k_ptr, const void* v_ptr, + void* o_ptr, void* softmax_lse_ptr, + void* softmax_lse_accum_ptr, void* o_accum_ptr, + int batch, int seqlen_q, int seqlen_k, + int num_heads_q, int num_heads_kv, int head_dim, + int q_batch_stride, int q_row_stride, int q_head_stride, + int k_batch_stride, int k_row_stride, int k_head_stride, + int v_batch_stride, int v_row_stride, int v_head_stride, + int o_batch_stride, int o_row_stride, int o_head_stride, + float softmax_scale, int num_sms, cudaStream_t stream); + +void fvk_attention_fa2_fwd_bf16_seqused( + const void* q_ptr, const void* k_ptr, const void* v_ptr, + void* o_ptr, void* softmax_lse_ptr, const void* seqused_k_ptr, + int batch, int seqlen_q, int seqlen_k, + int num_heads_q, int num_heads_kv, int head_dim, + int q_batch_stride, int q_row_stride, int q_head_stride, + int k_batch_stride, int k_row_stride, int k_head_stride, + int v_batch_stride, int v_row_stride, int v_head_stride, + int o_batch_stride, int o_row_stride, int o_head_stride, + float softmax_scale, int num_sms, cudaStream_t stream); + +void fvk_attention_fa2_fwd_bf16_seqused_splitkv( + const void* q_ptr, const void* k_ptr, const void* v_ptr, + void* o_ptr, void* softmax_lse_ptr, const void* seqused_k_ptr, + void* softmax_lse_accum_ptr, void* o_accum_ptr, + int batch, int seqlen_q, int seqlen_k, + int num_heads_q, int num_heads_kv, int head_dim, + int q_batch_stride, int q_row_stride, int q_head_stride, + int k_batch_stride, int k_row_stride, int k_head_stride, + int v_batch_stride, int v_row_stride, int v_head_stride, + int o_batch_stride, int o_row_stride, int o_head_stride, + float softmax_scale, int num_sms, cudaStream_t stream); + +void fvk_attention_fa2_fwd_bf16_causal( + const void* q_ptr, const void* k_ptr, const void* v_ptr, + void* o_ptr, void* softmax_lse_ptr, + void* softmax_lse_accum_ptr, void* o_accum_ptr, + int batch, int seqlen_q, int seqlen_k, + int num_heads_q, int num_heads_kv, int head_dim, + int q_batch_stride, int q_row_stride, int q_head_stride, + int k_batch_stride, int k_row_stride, int k_head_stride, + int v_batch_stride, int v_row_stride, int v_head_stride, + int o_batch_stride, int o_row_stride, int o_head_stride, + float softmax_scale, int num_sms, cudaStream_t stream); + +#ifdef __cplusplus +} +#endif + +#endif // FLASHRT_ATTENTION_FA2_WRAPPER_H diff --git a/csrc/attention/fa2_wrapper_causal.cu b/csrc/attention/fa2_wrapper_causal.cu index f651b194..fda63187 100644 --- a/csrc/attention/fa2_wrapper_causal.cu +++ b/csrc/attention/fa2_wrapper_causal.cu @@ -25,6 +25,7 @@ #include "flash_attn_2_src/flash_attn/namespace_config.h" #include "flash_attn_2_src/flash_attn/flash.h" +#include "attention/fa2_wrapper.h" namespace FLASH_NAMESPACE { template @@ -180,20 +181,17 @@ extern "C" void fvk_attention_fa2_fwd_bf16_causal( int o_batch_stride, int o_row_stride, int o_head_stride, float softmax_scale, int num_sms, cudaStream_t stream) { - if ((head_dim != 128) -#ifdef FA2_HAS_HDIM_256 - && head_dim != 256 + bool supported = false; +#if defined(FA2_HAS_BF16) && defined(FA2_HAS_HDIM_128) + supported = supported || head_dim == 128; #endif - ) { -#ifdef FA2_HAS_HDIM_256 - fprintf(stderr, - "fvk_attention_fa2_fwd_bf16_causal: head_dim=%d not built. " - "Only head_dim=128 and 256 are currently instantiated.\n", head_dim); -#else +#if defined(FA2_HAS_BF16) && defined(FA2_HAS_HDIM_256) + supported = supported || head_dim == 256; +#endif + if (!supported) { fprintf(stderr, "fvk_attention_fa2_fwd_bf16_causal: head_dim=%d not built. " - "Only head_dim=128 is currently instantiated.\n", head_dim); -#endif + "Enable its FA2_HDIMS entry and rebuild.\n", head_dim); std::abort(); } @@ -208,26 +206,36 @@ extern "C" void fvk_attention_fa2_fwd_bf16_causal( o_batch_stride, o_row_stride, o_head_stride, softmax_scale); - int num_splits = setup_splitkv_causal(params, softmax_lse_accum_ptr, o_accum_ptr, - num_sms, seqlen_q, seqlen_k, - head_dim, batch, num_heads_q); - if (head_dim == 128 && num_splits > 1) { - FLASH_NAMESPACE::run_mha_fwd_splitkv_dispatch(params, stream); - } else if (head_dim == 128) { - FLASH_NAMESPACE::run_mha_fwd_(params, stream); - } -#ifdef FA2_HAS_HDIM_256 - else if (num_splits > 1) { - FLASH_NAMESPACE::run_mha_fwd_splitkv_dispatch(params, stream); - } else { - FLASH_NAMESPACE::run_mha_fwd_(params, stream); - } -#else - else { - fprintf(stderr, - "fvk_attention_fa2_fwd_bf16_causal: head_dim=%d not built " - "(hdim=256 disabled at compile time).\n", head_dim); - std::abort(); - } + int num_splits = setup_splitkv_causal( + params, softmax_lse_accum_ptr, o_accum_ptr, num_sms, seqlen_q, + seqlen_k, head_dim, batch, num_heads_q); + switch (head_dim) { +#if defined(FA2_HAS_BF16) && defined(FA2_HAS_HDIM_128) + case 128: + if (num_splits > 1) { + FLASH_NAMESPACE::run_mha_fwd_splitkv_dispatch< + cutlass::bfloat16_t, 128, true>(params, stream); + } else { + FLASH_NAMESPACE::run_mha_fwd_< + cutlass::bfloat16_t, 128, true>(params, stream); + } + return; +#endif +#if defined(FA2_HAS_BF16) && defined(FA2_HAS_HDIM_256) + case 256: + if (num_splits > 1) { + FLASH_NAMESPACE::run_mha_fwd_splitkv_dispatch< + cutlass::bfloat16_t, 256, true>(params, stream); + } else { + FLASH_NAMESPACE::run_mha_fwd_< + cutlass::bfloat16_t, 256, true>(params, stream); + } + return; #endif + default: + fprintf(stderr, + "fvk_attention_fa2_fwd_bf16_causal: head_dim=%d not built " + "in this FA2 matrix.\n", head_dim); + std::abort(); + } } diff --git a/csrc/fa2_bindings.cpp b/csrc/fa2_bindings.cpp index 41d13d04..75360388 100644 --- a/csrc/fa2_bindings.cpp +++ b/csrc/fa2_bindings.cpp @@ -22,79 +22,14 @@ #include #include +#include "attention/fa2_wrapper.h" + namespace py = pybind11; static cudaStream_t to_stream(uintptr_t s) { return reinterpret_cast(s); } -// Forward declarations (definitions in csrc/attention/fa2_wrapper.cu). -extern "C" void fvk_attention_fa2_fwd_fp16( - const void* q_ptr, const void* k_ptr, const void* v_ptr, - void* o_ptr, void* softmax_lse_ptr, - void* softmax_lse_accum_ptr, void* o_accum_ptr, - int batch, int seqlen_q, int seqlen_k, - int num_heads_q, int num_heads_kv, int head_dim, - int q_batch_stride, int q_row_stride, int q_head_stride, - int k_batch_stride, int k_row_stride, int k_head_stride, - int v_batch_stride, int v_row_stride, int v_head_stride, - int o_batch_stride, int o_row_stride, int o_head_stride, - float softmax_scale, int num_sms, cudaStream_t stream); - -extern "C" void fvk_attention_fa2_fwd_bf16( - const void* q_ptr, const void* k_ptr, const void* v_ptr, - void* o_ptr, void* softmax_lse_ptr, - void* softmax_lse_accum_ptr, void* o_accum_ptr, - int batch, int seqlen_q, int seqlen_k, - int num_heads_q, int num_heads_kv, int head_dim, - int q_batch_stride, int q_row_stride, int q_head_stride, - int k_batch_stride, int k_row_stride, int k_head_stride, - int v_batch_stride, int v_row_stride, int v_head_stride, - int o_batch_stride, int o_row_stride, int o_head_stride, - float softmax_scale, int num_sms, cudaStream_t stream); - -// seqused_k variant — definition in csrc/attention/fa2_wrapper.cu. Reads the -// per-batch K length from device memory so one captured graph serves any pos. -extern "C" void fvk_attention_fa2_fwd_bf16_seqused( - const void* q_ptr, const void* k_ptr, const void* v_ptr, - void* o_ptr, void* softmax_lse_ptr, const void* seqused_k_ptr, - int batch, int seqlen_q, int seqlen_k, - int num_heads_q, int num_heads_kv, int head_dim, - int q_batch_stride, int q_row_stride, int q_head_stride, - int k_batch_stride, int k_row_stride, int k_head_stride, - int v_batch_stride, int v_row_stride, int v_head_stride, - int o_batch_stride, int o_row_stride, int o_head_stride, - float softmax_scale, int num_sms, cudaStream_t stream); - -// seqused_k + split-KV variant (experimental; caller pre-inits lse_accum=-inf). -extern "C" void fvk_attention_fa2_fwd_bf16_seqused_splitkv( - const void* q_ptr, const void* k_ptr, const void* v_ptr, - void* o_ptr, void* softmax_lse_ptr, const void* seqused_k_ptr, - void* softmax_lse_accum_ptr, void* o_accum_ptr, - int batch, int seqlen_q, int seqlen_k, - int num_heads_q, int num_heads_kv, int head_dim, - int q_batch_stride, int q_row_stride, int q_head_stride, - int k_batch_stride, int k_row_stride, int k_head_stride, - int v_batch_stride, int v_row_stride, int v_head_stride, - int o_batch_stride, int o_row_stride, int o_head_stride, - float softmax_scale, int num_sms, cudaStream_t stream); - -// Causal variant — definition in csrc/attention/fa2_wrapper_causal.cu. -// Currently only (bf16, head_dim=128) is built. Used by Qwen3-8B -// prefill (S=N causal self-attention). -extern "C" void fvk_attention_fa2_fwd_bf16_causal( - const void* q_ptr, const void* k_ptr, const void* v_ptr, - void* o_ptr, void* softmax_lse_ptr, - void* softmax_lse_accum_ptr, void* o_accum_ptr, - int batch, int seqlen_q, int seqlen_k, - int num_heads_q, int num_heads_kv, int head_dim, - int q_batch_stride, int q_row_stride, int q_head_stride, - int k_batch_stride, int k_row_stride, int k_head_stride, - int v_batch_stride, int v_row_stride, int v_head_stride, - int o_batch_stride, int o_row_stride, int o_head_stride, - float softmax_scale, int num_sms, cudaStream_t stream); - - // Shared docstring. pybind::def's doc arg takes a single string; we want the // same text for both fwd_fp16 and fwd_bf16 so deduplicate via static const. static const char* kDocstring = R"(FlashAttention-2 fwd (vendored). GQA-capable cross-attention. diff --git a/docs/fa2_native_c_api.md b/docs/fa2_native_c_api.md new file mode 100644 index 00000000..2d7816c1 --- /dev/null +++ b/docs/fa2_native_c_api.md @@ -0,0 +1,46 @@ +# FA2 native C library + +FlashRT builds vendored FA2 entry points into a Python-free shared library: + +- `libflashrt_fa2_raw.so` owns the `fvk_attention_fa2_fwd_*` C symbols; +- `flash_rt_fa2` is the Python adapter and links the raw library; +- native C++ runtimes link the same raw library directly. + +This keeps one implementation and one symbol set for Python and native +producers. Model-specific code must not compile a private copy of the FA2 +wrappers. + +This is an internal native linkage boundary, not a new versioned public ABI. +Consumers outside the FlashRT build must continue to use the supported +runtime interfaces rather than bind these kernel entry points directly. + +## Packaging contract + +The Python adapter and raw library are one install unit. Deployments that use +`flash_rt_fa2` must install both files in the same directory. Both targets use +an `$ORIGIN` runtime search path, so no build-tree path is part of deployment. + +Copying only the Python extension is no longer supported. Packaging and image +rules should select the CMake install targets instead of copying an individual +extension file. + +## C boundary + +The raw library has no Python dependency. Its declarations live in +`csrc/attention/fa2_wrapper.h`; the Python adapter includes that header instead +of maintaining a second set of declarations. Existing `fvk_*` signatures are +unchanged. + +The causal entry is always present when FA2 is enabled. Unsupported dtype/head +dimension combinations fail explicitly according to the compiled FA2 matrix; +adding a model consumer must not silently change that matrix. + +## Validation + +- the raw library exports the expected `fvk_attention_fa2_fwd_*` symbols; +- it has no unresolved Python or `fvk_*` symbols; +- the Python adapter has a dynamic dependency on the raw library; +- adapter and raw-library runtime search paths are `$ORIGIN`; +- the install smoke test rejects an adapter without its raw library; +- the Python adapter and native C++ attention gates both execute against the + shared implementation. diff --git a/setup.py b/setup.py index 6b7453ea..6d65f979 100644 --- a/setup.py +++ b/setup.py @@ -21,7 +21,8 @@ cmake --build build -j After this, ``flash_rt/flash_rt_kernels*.so`` (and on RTX, -``flash_rt_fa2*.so``; on Thor/Hopper, ``flash_rt_fp4*.so``) exist +``flash_rt_fa2*.so`` plus ``libflashrt_fa2_raw.so``; on Thor/Hopper, +``flash_rt_fp4*.so``) exist and ``import flash_rt`` works in editable installs. Optional pip dependency: the legacy upstream attention path diff --git a/tests/test_install_smoke.py b/tests/test_install_smoke.py index b931f082..e692397e 100644 --- a/tests/test_install_smoke.py +++ b/tests/test_install_smoke.py @@ -18,7 +18,8 @@ torch 2.5.1 / CUDA 12.5 / Py3.12 → 30–60 min sdist compile per cold image). The vendored ``flash_rt_fa2.so`` is enough for the default path; ``flash-attn`` is only needed for legacy bisection - sites. + sites. The vendored adapter and ``libflashrt_fa2_raw.so`` form one + install unit. Run: PYTHONPATH=. python -m pytest tests/test_install_smoke.py -v @@ -33,6 +34,20 @@ import pytest +def test_fa2_install_unit_is_complete(): + """A built FA2 adapter must be deployed with its Python-free raw library.""" + import flash_rt + + pkg_dir = pathlib.Path(flash_rt.__file__).parent + adapters = list(pkg_dir.glob("flash_rt_fa2*.so")) \ + + list(pkg_dir.glob("flash_rt_fa2*.pyd")) + if not adapters: + pytest.skip("vendored FA2 is not configured in this build") + raw = list(pkg_dir.glob("libflashrt_fa2_raw.so*")) \ + + list(pkg_dir.glob("flashrt_fa2_raw*.dll")) + assert raw, "flash_rt_fa2 was built without libflashrt_fa2_raw" + + def test_kernels_so_lives_in_flash_rt(): """The compiled .so must land inside the ``flash_rt/`` package at build time. ``cmake --build`` alone should be sufficient — no