refactor(fa2): expose a Python-free C runtime

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,19 +954,39 @@ 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_OBJECTS:fa2_vendor_obj>)
+  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(...)
   pybind11_add_module(flash_rt_fa2 csrc/fa2_bindings.cpp)
   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_OBJECTS:fa2_vendor_obj>)
   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.)
 

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<typename elem_type, int kHeadDim, bool Is_causal>

csrc/attention/fa2_wrapper.h

@@ -0,0 +1,73 @@
+#ifndef FLASHRT_ATTENTION_FA2_WRAPPER_H
+#define FLASHRT_ATTENTION_FA2_WRAPPER_H
+
+#include <cuda_runtime_api.h>
+
+#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

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<typename elem_type, int kHeadDim, bool Is_causal>
@@ -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<cutlass::bfloat16_t, 128, true>(params, stream);
-    } else if (head_dim == 128) {
-        FLASH_NAMESPACE::run_mha_fwd_<cutlass::bfloat16_t, 128, true>(params, stream);
-    }
-#ifdef FA2_HAS_HDIM_256
-    else 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);
-    }
-#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();
+    }
 }

csrc/fa2_bindings.cpp

@@ -22,79 +22,14 @@
 #include <cuda_runtime.h>
 #include <cstdint>
 
+#include "attention/fa2_wrapper.h"
+
 namespace py = pybind11;
 
 static cudaStream_t to_stream(uintptr_t s) {
     return reinterpret_cast<cudaStream_t>(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.

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.

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