feat: execute FlashKDA on RTX 5090

experiments/k3/flashkda/build_reference_helper.py

@@ -0,0 +1,50 @@
+#!/usr/bin/env python3
+"""Prebuild the CUDA helper used verbatim by upstream tests/torch_ref.py.
+
+This runs in the CUDA build container without a GPU. Keeping the extension
+name, sources, generated function, and compiler flags identical lets PyTorch
+reuse the cache when the official reference module is imported on the host.
+"""
+
+from torch.utils.cpp_extension import load_inline
+
+
+CUDA_SOURCE = r"""
+#include <torch/extension.h>
+#include <cuda_runtime.h>
+
+__global__ void sigmoid_tanh_fp32_kernel(const float* __restrict__ input,
+                                         float* __restrict__ output, int n) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx < n) {
+        float xh = input[idx] * 0.5f;
+        float th;
+        asm("tanh.approx.f32 %0, %1;" : "=f"(th) : "f"(xh));
+        output[idx] = th * 0.5f + 0.5f;
+    }
+}
+
+torch::Tensor sigmoid_tanh_fp32(torch::Tensor input) {
+    auto output = torch::empty_like(input);
+    int n = input.numel();
+    sigmoid_tanh_fp32_kernel<<<(n + 255) / 256, 256>>>(
+        input.data_ptr<float>(), output.data_ptr<float>(), n);
+    return output;
+}
+"""
+
+
+def main() -> None:
+    module = load_inline(
+        name="sigmoid_ext",
+        cpp_sources="torch::Tensor sigmoid_tanh_fp32(torch::Tensor input);",
+        cuda_sources=CUDA_SOURCE,
+        functions=["sigmoid_tanh_fp32"],
+        extra_cuda_cflags=["-O2"],
+        verbose=True,
+    )
+    print(module.__file__)
+
+
+if __name__ == "__main__":
+    main()