From 663dfb6f19abbc8e8de26dfac7f9634ac2ece0e9 Mon Sep 17 00:00:00 2001 From: Giorgis Georgakoudis Date: Fri, 13 Mar 2026 07:01:17 -0700 Subject: [PATCH 1/2] Update for proteus PR https://github.com/Olympus-HPC/proteus/pull/429 - CppJitModule::instantiate returns a reference --- csrc/flash_attn/src/flash_bwd_jit_bridge.cu | 6 +++--- csrc/flash_attn/src/flash_fwd_jit_bridge.cu | 4 ++-- csrc/flash_attn/src/flash_fwd_splitkv_jit_bridge.cu | 4 ++-- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/csrc/flash_attn/src/flash_bwd_jit_bridge.cu b/csrc/flash_attn/src/flash_bwd_jit_bridge.cu index 5565d465342..586356b6072 100644 --- a/csrc/flash_attn/src/flash_bwd_jit_bridge.cu +++ b/csrc/flash_attn/src/flash_bwd_jit_bridge.cu @@ -540,7 +540,7 @@ bool try_run_mha_bwd_jit_impl(Flash_bwd_params ¶ms, cudaStream_t stream, // Kernel 1: dot_do_o (preprocess) — smem = 0 // ----------------------------------------------------------------------- { - auto Instance = get_bwd_module().instantiate( + auto &Instance = get_bwd_module().instantiate( "flash_bwd_dot_do_o_jit", to_template_bool(!params.deterministic), // Clear_dQaccum Config->KernelTraitsName); @@ -554,7 +554,7 @@ bool try_run_mha_bwd_jit_impl(Flash_bwd_params ¶ms, cudaStream_t stream, // ----------------------------------------------------------------------- { const size_t SmemSize = Config->SmemSize1colblock; - auto Instance = get_bwd_module().instantiate( + auto &Instance = get_bwd_module().instantiate( "flash_bwd_dq_dk_dv_jit", Config->KernelTraitsName, to_template_bool(Is_dropout_t), to_template_bool(Is_causal), to_template_bool(Is_local_t), to_template_bool(Has_alibi), @@ -571,7 +571,7 @@ bool try_run_mha_bwd_jit_impl(Flash_bwd_params ¶ms, cudaStream_t stream, { const size_t SmemSize = Config->SmemdQSize; const int nsplits = !params.deterministic ? 1 : gridDimx; - auto Instance = get_bwd_module().instantiate( + auto &Instance = get_bwd_module().instantiate( "flash_bwd_convert_dq_jit", Config->KernelTraitsName); Instance.launch(GridM, Block, SmemSize, reinterpret_cast(stream), params, nsplits); diff --git a/csrc/flash_attn/src/flash_fwd_jit_bridge.cu b/csrc/flash_attn/src/flash_fwd_jit_bridge.cu index 4251564c886..ed62238be76 100644 --- a/csrc/flash_attn/src/flash_fwd_jit_bridge.cu +++ b/csrc/flash_attn/src/flash_fwd_jit_bridge.cu @@ -454,7 +454,7 @@ bool try_run_mha_fwd_jit_fp16(Flash_fwd_params ¶ms, cudaStream_t stream) { const dim3 Block(Config->NThreads, 1, 1); const size_t SmemSize = Config->SmemSize; - auto Instance = get_module().instantiate( + auto &Instance = get_module().instantiate( "flash_fwd_kernel_jit", Config->KernelTraitsName, to_template_bool(Is_dropout_t), to_template_bool(Is_causal), to_template_bool(Is_local), to_template_bool(Has_alibi), @@ -510,7 +510,7 @@ bool try_run_mha_fwd_jit_bf16(Flash_fwd_params ¶ms, cudaStream_t stream) { const dim3 Block(Config->NThreads, 1, 1); const size_t SmemSize = Config->SmemSize; - auto Instance = get_module().instantiate( + auto &Instance = get_module().instantiate( "flash_fwd_kernel_jit", Config->KernelTraitsName, to_template_bool(Is_dropout_t), to_template_bool(Is_causal), to_template_bool(Is_local), to_template_bool(Has_alibi), diff --git a/csrc/flash_attn/src/flash_fwd_splitkv_jit_bridge.cu b/csrc/flash_attn/src/flash_fwd_splitkv_jit_bridge.cu index 5a1f1be2dc8..7e4ea592f60 100644 --- a/csrc/flash_attn/src/flash_fwd_splitkv_jit_bridge.cu +++ b/csrc/flash_attn/src/flash_fwd_splitkv_jit_bridge.cu @@ -297,7 +297,7 @@ bool try_run_mha_fwd_splitkv_jit_impl( const dim3 Block(Config->NThreads, 1, 1); const size_t SmemSize = Config->SmemSize; - auto MainInstance = get_splitkv_module().instantiate( + auto &MainInstance = get_splitkv_module().instantiate( "flash_fwd_splitkv_kernel_jit", Config->KernelTraitsName, to_template_bool(Is_causal), to_template_bool(Is_local && !Is_causal), @@ -334,7 +334,7 @@ bool try_run_mha_fwd_splitkv_jit_impl( (params.b * params.h * params.seqlen_q + CombBlockM - 1) / CombBlockM); const dim3 BlockCombine(Config->NThreads, 1, 1); - auto CombineInstance = get_splitkv_module().instantiate( + auto &CombineInstance = get_splitkv_module().instantiate( "flash_fwd_splitkv_combine_kernel_jit", Config->KernelTraitsName, std::to_string(CombBlockM).c_str(), std::to_string(LogMaxSplits).c_str(), From 711b1700f2321ba8195ee7a1a5f2b59548263e8c Mon Sep 17 00:00:00 2001 From: Giorgis Georgakoudis Date: Thu, 19 Mar 2026 19:15:05 -0700 Subject: [PATCH 2/2] Updates --- benchmarks/benchmark_flash_attention.py | 34 +++++++----- build_fa_proteus.sh | 16 +++--- csrc/flash_attn/src/flash_bwd_jit_bridge.cu | 54 ++++++++++++++----- csrc/flash_attn/src/flash_fwd_jit_bridge.cu | 41 ++++++++++---- .../src/flash_fwd_splitkv_jit_bridge.cu | 24 +++++++-- setup.py | 1 + 6 files changed, 126 insertions(+), 44 deletions(-) diff --git a/benchmarks/benchmark_flash_attention.py b/benchmarks/benchmark_flash_attention.py index 73ec9e28eb5..2a944a46750 100644 --- a/benchmarks/benchmark_flash_attention.py +++ b/benchmarks/benchmark_flash_attention.py @@ -4,16 +4,18 @@ import math import os -if os.getenv("FLASH_ATTN_USE_PROTEUS_JIT") == "1": - # PyTorch's Proteus softmax path reads this at module-init time, so ensure - # it exists before importing torch. - cwd_softmax = "smax_kernel.cu" - local_softmax = os.path.join(os.path.dirname(__file__), "smax_kernel.cu") - if (not os.path.exists(cwd_softmax)) and os.path.exists(local_softmax): - try: - os.symlink(local_softmax, cwd_softmax) - except FileExistsError: - pass +# This is needed only if you use the modified PyTorch for proteus compilation +# and you do a pytorch execution of the benchmark. +#if os.getenv("FLASH_ATTN_USE_PROTEUS_JIT") == "1": +# # PyTorch's Proteus softmax path reads this at module-init time, so ensure +# # it exists before importing torch. +# cwd_softmax = "smax_kernel.cu" +# local_softmax = os.path.join(os.path.dirname(__file__), "smax_kernel.cu") +# if (not os.path.exists(cwd_softmax)) and os.path.exists(local_softmax): +# try: +# os.symlink(local_softmax, cwd_softmax) +# except FileExistsError: +# pass import torch import torch.nn as nn @@ -90,10 +92,14 @@ def time_fwd_bwd(func, *args, **kwargs): device = 'cuda' dtype = torch.float16 -bs_seqlen_vals = [(32, 512), (16, 1024), (8, 2048), (4, 4096), (2, 8192), (1, 16384)] -causal_vals = [False, True] -headdim_vals = [64, 128] -dtype_vals = [torch.float16, torch.bfloat16] +#bs_seqlen_vals = [(32, 512), (16, 1024), (8, 2048), (4, 4096), (2, 8192), (1, 16384)] +bs_seqlen_vals = [(32, 512)] +#causal_vals = [False, True] +causal_vals = [False] +#headdim_vals = [64, 128] +headdim_vals = [64] +#dtype_vals = [torch.float16, torch.bfloat16] +dtype_vals = [torch.float16] dim = 2048 dropout_p = 0.0 diff --git a/build_fa_proteus.sh b/build_fa_proteus.sh index d8a94563599..9a27949460a 100755 --- a/build_fa_proteus.sh +++ b/build_fa_proteus.sh @@ -53,7 +53,7 @@ PROTEUS_BRANCH="${PROTEUS_BRANCH:-main}" PROTEUS_INSTALL="${PROTEUS_INSTALL:-}" # CUDA architectures to compile for (semicolon-separated) -CUDA_ARCHS="${CUDA_ARCHS:-80;90}" +CUDA_ARCHS="${CUDA_ARCHS:-90}" # Number of parallel build jobs JOBS="${JOBS:-$(nproc)}" @@ -117,7 +117,7 @@ else git clone "${PROTEUS_REPO}" "${PROTEUS_SRC}" fi cd "${PROTEUS_SRC}" - git checkout "${PROTEUS_BRANCH}" + #git checkout "${PROTEUS_BRANCH}" cd "${INSTALL_DIR}" mkdir -p "${PROTEUS_BUILD}" @@ -133,13 +133,15 @@ else -DCMAKE_CXX_COMPILER="${CLANG_BIN}/clang++" \ -DCMAKE_INSTALL_PREFIX="${PROTEUS_INSTALL}" \ -DCMAKE_BUILD_TYPE=RelWithDebInfo \ + -DCMAKE_CUDA_ARCHITECTURES="${CUDA_ARCHS}" \ + -DCMAKE_CUDA_COMPILER="${CLANG_BIN}/clang++" \ -DCMAKE_EXPORT_COMPILE_COMMANDS=ON echo "Building Proteus (${JOBS} jobs)..." make -j"${JOBS}" echo "Installing Proteus..." - make install + make -j"${JOBS}" install echo "Proteus installed to: ${PROTEUS_INSTALL}" cd "${INSTALL_DIR}" @@ -164,8 +166,10 @@ echo "=== Step 2: Building Flash-Attention with Proteus JIT ===" FA_SRC="${INSTALL_DIR}/flash-attention" if [ ! -d "${FA_SRC}" ]; then - echo "Cloning Flash-Attention (with submodules)..." - git clone --recurse-submodules -b "${FA_BRANCH}" "${FA_REPO}" "${FA_SRC}" + echo "Missing Flash-Attention source at ${FA_SRC}" + exit 1 + #echo "Cloning Flash-Attention (with submodules)..." + #git clone --recurse-submodules -b "${FA_BRANCH}" "${FA_REPO}" "${FA_SRC}" fi cd "${FA_SRC}" @@ -177,7 +181,7 @@ PROTEUS_PREFIX="${PROTEUS_INSTALL}" \ PROTEUS_LLVM_LIBDIR="${LLVM_DIR}" \ PROTEUS_CLANG_LIBDIR="${LLVM_DIR}" \ PROTEUS_CUDA_ROOT="${CUDA_ROOT}" \ -"${PYTHON_BIN}" setup.py build_ext --inplace +"${PYTHON_BIN}" setup.py build_ext --build-lib build/fa-proteus --build-temp build/fa-proteus/temp # Verify the .so was built FA_SO=$(ls "${FA_SRC}"/flash_attn_2_cuda*.so 2>/dev/null | head -1) diff --git a/csrc/flash_attn/src/flash_bwd_jit_bridge.cu b/csrc/flash_attn/src/flash_bwd_jit_bridge.cu index 586356b6072..029ac932a0e 100644 --- a/csrc/flash_attn/src/flash_bwd_jit_bridge.cu +++ b/csrc/flash_attn/src/flash_bwd_jit_bridge.cu @@ -218,7 +218,8 @@ bool use_proteus_jit() { return Enabled; } -std::vector get_extra_args() { +std::vector get_extra_args( + proteus::CppJitCompilerBackend Backend) { const std::string SourceRoot = get_source_root(); const std::string TorchInclude = get_torch_include_dir(SourceRoot); const std::string TorchApiInclude = get_torch_api_include_dir(TorchInclude); @@ -231,7 +232,16 @@ std::vector get_extra_args() { "-I" + TorchApiInclude, "-I" + TorchAtenInclude, "-I" + CudaInclude, + "-U__CUDA_NO_HALF_OPERATORS__", + "-U__CUDA_NO_HALF_CONVERSIONS__", + "-U__CUDA_NO_HALF2_OPERATORS__", + "-U__CUDA_NO_BFLOAT16_CONVERSIONS__", }; + if (Backend == proteus::CppJitCompilerBackend::Nvcc) { + Args.push_back("--expt-relaxed-constexpr"); + Args.push_back("--expt-extended-lambda"); + Args.push_back("--use_fast_math"); + } if (const char *DebugArgs = std::getenv("FLASH_ATTN_JIT_DEBUG_ARGS"); DebugArgs != nullptr && std::strcmp(DebugArgs, "1") == 0) { for (const auto &Arg : Args) { @@ -366,8 +376,9 @@ __global__ void flash_bwd_convert_dq_jit(const flash_bwd_params_t params, auto &get_bwd_module() { static const std::string JitKernelCode = get_bwd_jit_kernel_code(); + static constexpr auto Backend = proteus::CppJitCompilerBackend::Nvcc; static auto Module = std::make_unique( - "cuda", JitKernelCode, get_extra_args()); + "cuda", JitKernelCode, get_extra_args(Backend), Backend); return *Module; } @@ -544,8 +555,14 @@ bool try_run_mha_bwd_jit_impl(Flash_bwd_params ¶ms, cudaStream_t stream, "flash_bwd_dot_do_o_jit", to_template_bool(!params.deterministic), // Clear_dQaccum Config->KernelTraitsName); - Instance.launch(GridM, Block, /*smem=*/0, - reinterpret_cast(stream), params); + auto Ret = Instance.launch(GridM, Block, /*smem=*/0, + reinterpret_cast(stream), params); + if (Ret != cudaSuccess) { + std::fprintf(stderr, "%s:%d: Kernel 1 launch failed with error code %s\n", + __FILE__, __LINE__, + cudaGetErrorString(static_cast(static_cast(Ret)))); + std::abort(); + } C10_CUDA_KERNEL_LAUNCH_CHECK(); } @@ -560,8 +577,16 @@ bool try_run_mha_bwd_jit_impl(Flash_bwd_params ¶ms, cudaStream_t stream, to_template_bool(Is_local_t), to_template_bool(Has_alibi), to_template_bool(Is_even_MN_t), to_template_bool(Is_even_K_t), to_template_bool(Is_softcap)); - Instance.launch(GridN, Block, SmemSize, - reinterpret_cast(stream), params); + Instance.setFuncAttribute( + proteus::CppJitFuncAttribute::MaxDynamicSharedMemorySize, SmemSize); + auto Ret = Instance.launch(GridN, Block, SmemSize, + reinterpret_cast(stream), params); + if (Ret != cudaSuccess) { + std::fprintf(stderr, "%s:%d: Kernel 1 launch failed with error code %s\n", + __FILE__, __LINE__, + cudaGetErrorString(static_cast(static_cast(Ret)))); + std::abort(); + } C10_CUDA_KERNEL_LAUNCH_CHECK(); } @@ -573,8 +598,17 @@ bool try_run_mha_bwd_jit_impl(Flash_bwd_params ¶ms, cudaStream_t stream, const int nsplits = !params.deterministic ? 1 : gridDimx; auto &Instance = get_bwd_module().instantiate( "flash_bwd_convert_dq_jit", Config->KernelTraitsName); - Instance.launch(GridM, Block, SmemSize, - reinterpret_cast(stream), params, nsplits); + Instance.setFuncAttribute( + proteus::CppJitFuncAttribute::MaxDynamicSharedMemorySize, SmemSize); + auto Ret = + Instance.launch(GridM, Block, SmemSize, + reinterpret_cast(stream), params, nsplits); + if (Ret != cudaSuccess) { + std::fprintf(stderr, "%s:%d: Kernel 1 launch failed with error code %s\n", + __FILE__, __LINE__, + cudaGetErrorString(static_cast(static_cast(Ret)))); + std::abort(); + } C10_CUDA_KERNEL_LAUNCH_CHECK(); } @@ -596,8 +630,6 @@ bool try_run_mha_bwd_jit_fp16(Flash_bwd_params ¶ms, cudaStream_t stream) { bool result = try_run_mha_bwd_jit_impl(params, stream, Config); auto t1 = std::chrono::steady_clock::now(); auto ms = std::chrono::duration_cast(t1 - t0).count(); - std::fprintf(stderr, "[proteus-jit] bwd_fp16 h=%d launch %ld ms\n", - Config ? Config->HeadDim : params.d, ms); return result; } @@ -610,8 +642,6 @@ bool try_run_mha_bwd_jit_bf16(Flash_bwd_params ¶ms, cudaStream_t stream) { bool result = try_run_mha_bwd_jit_impl(params, stream, Config); auto t1 = std::chrono::steady_clock::now(); auto ms = std::chrono::duration_cast(t1 - t0).count(); - std::fprintf(stderr, "[proteus-jit] bwd_bf16 h=%d launch %ld ms\n", - Config ? Config->HeadDim : params.d, ms); return result; } diff --git a/csrc/flash_attn/src/flash_fwd_jit_bridge.cu b/csrc/flash_attn/src/flash_fwd_jit_bridge.cu index ed62238be76..21266393c6b 100644 --- a/csrc/flash_attn/src/flash_fwd_jit_bridge.cu +++ b/csrc/flash_attn/src/flash_fwd_jit_bridge.cu @@ -236,7 +236,8 @@ bool use_proteus_jit() { return Enabled; } -std::vector get_extra_args() { +std::vector get_extra_args( + proteus::CppJitCompilerBackend Backend) { const std::string SourceRoot = get_source_root(); const std::string TorchInclude = get_torch_include_dir(SourceRoot); const std::string TorchApiInclude = get_torch_api_include_dir(TorchInclude); @@ -249,7 +250,16 @@ std::vector get_extra_args() { "-I" + TorchApiInclude, "-I" + TorchAtenInclude, "-I" + CudaInclude, + "-U__CUDA_NO_HALF_OPERATORS__", + "-U__CUDA_NO_HALF_CONVERSIONS__", + "-U__CUDA_NO_HALF2_OPERATORS__", + "-U__CUDA_NO_BFLOAT16_CONVERSIONS__", }; + if (Backend == proteus::CppJitCompilerBackend::Nvcc) { + Args.push_back("--expt-relaxed-constexpr"); + Args.push_back("--expt-extended-lambda"); + Args.push_back("--use_fast_math"); + } if (const char *DebugArgs = std::getenv("FLASH_ATTN_JIT_DEBUG_ARGS"); DebugArgs != nullptr && std::strcmp(DebugArgs, "1") == 0) { for (const auto &Arg : Args) { @@ -261,8 +271,9 @@ std::vector get_extra_args() { auto &get_module() { static const std::string JitKernelCode = get_jit_kernel_code(); + static constexpr auto Backend = proteus::CppJitCompilerBackend::Nvcc; static auto Module = std::make_unique( - "cuda", JitKernelCode, get_extra_args()); + "cuda", JitKernelCode, get_extra_args(Backend), Backend); return *Module; } @@ -460,14 +471,21 @@ bool try_run_mha_fwd_jit_fp16(Flash_fwd_params ¶ms, cudaStream_t stream) { to_template_bool(Is_local), to_template_bool(Has_alibi), to_template_bool(Is_even_MN_t), to_template_bool(Is_even_K_t), to_template_bool(Is_softcap), to_template_bool(Return_softmax_t)); + Instance.setFuncAttribute(proteus::CppJitFuncAttribute::MaxDynamicSharedMemorySize, SmemSize); auto t0 = std::chrono::steady_clock::now(); - Instance.launch(Grid, Block, SmemSize, reinterpret_cast(stream), - params); + auto Ret = Instance.launch(Grid, Block, SmemSize, + reinterpret_cast(stream), params); auto t1 = std::chrono::steady_clock::now(); + + if (Ret != cudaSuccess) { + std::fprintf(stderr, "%s:%d: Failed to launch flash_fwd_kernel_jit: %s\n", + __FILE__, __LINE__, + cudaGetErrorString(static_cast(static_cast(Ret)))); + std::abort(); + } + auto ms = std::chrono::duration_cast(t1 - t0).count(); - std::fprintf(stderr, "[proteus-jit] fwd_fp16 h=%d launch %ld ms\n", - Config->HeadDim, ms); C10_CUDA_KERNEL_LAUNCH_CHECK(); return true; } @@ -516,14 +534,19 @@ bool try_run_mha_fwd_jit_bf16(Flash_fwd_params ¶ms, cudaStream_t stream) { to_template_bool(Is_local), to_template_bool(Has_alibi), to_template_bool(Is_even_MN_t), to_template_bool(Is_even_K_t), to_template_bool(Is_softcap), to_template_bool(Return_softmax_t)); + Instance.setFuncAttribute(proteus::CppJitFuncAttribute::MaxDynamicSharedMemorySize, SmemSize); auto t0 = std::chrono::steady_clock::now(); - Instance.launch(Grid, Block, SmemSize, reinterpret_cast(stream), + auto Ret = Instance.launch(Grid, Block, SmemSize, reinterpret_cast(stream), params); auto t1 = std::chrono::steady_clock::now(); + if (static_cast(Ret) != cudaSuccess) { + std::fprintf(stderr, "%s:%d: Failed to launch flash_fwd_kernel_jit: %s\n", + __FILE__, __LINE__, + cudaGetErrorString(static_cast(static_cast(Ret)))); + std::abort(); + } auto ms = std::chrono::duration_cast(t1 - t0).count(); - std::fprintf(stderr, "[proteus-jit] fwd_bf16 h=%d launch %ld ms\n", - Config->HeadDim, ms); C10_CUDA_KERNEL_LAUNCH_CHECK(); return true; } diff --git a/csrc/flash_attn/src/flash_fwd_splitkv_jit_bridge.cu b/csrc/flash_attn/src/flash_fwd_splitkv_jit_bridge.cu index 7e4ea592f60..fefdb193ff2 100644 --- a/csrc/flash_attn/src/flash_fwd_splitkv_jit_bridge.cu +++ b/csrc/flash_attn/src/flash_fwd_splitkv_jit_bridge.cu @@ -227,7 +227,8 @@ __global__ void flash_fwd_splitkv_combine_kernel_jit(const flash_fwd_params_t pa auto &get_splitkv_module() { static const std::string Code = get_splitkv_jit_kernel_code(); static auto Module = std::make_unique( - "cuda", Code, get_extra_args()); + "cuda", Code, get_extra_args(), proteus::CppJitCompilerBackend::Nvcc); + return *Module; } @@ -308,10 +309,19 @@ bool try_run_mha_fwd_splitkv_jit_impl( to_template_bool(Split), to_template_bool(Append_KV)); + MainInstance.setFuncAttribute( + proteus::CppJitFuncAttribute::MaxDynamicSharedMemorySize, SmemSize); + auto t0_main = std::chrono::steady_clock::now(); - MainInstance.launch(Grid, Block, SmemSize, + auto Ret = MainInstance.launch(Grid, Block, SmemSize, reinterpret_cast(stream), params); auto t1_main = std::chrono::steady_clock::now(); + if (static_cast(Ret) != cudaSuccess) { + std::fprintf( + stderr, "%s:%d: Failed to launch flash_fwd_splitkv_kernel_jit: %s\n", + __FILE__, __LINE__, cudaGetErrorString(static_cast(static_cast(Ret)))); + std::abort(); + } auto ms_main = std::chrono::duration_cast(t1_main - t0_main).count(); std::fprintf(stderr, "[proteus-jit] %s h=%d main_launch %ld ms\n", DtypeLabel, Config->HeadDim, ms_main); @@ -341,10 +351,18 @@ bool try_run_mha_fwd_splitkv_jit_impl( to_template_bool(Is_even_K_t)); auto t0_comb = std::chrono::steady_clock::now(); - CombineInstance.launch(GridCombine, BlockCombine, /*smem=*/0, + auto Ret = CombineInstance.launch(GridCombine, BlockCombine, /*smem=*/0, reinterpret_cast(stream), params); auto t1_comb = std::chrono::steady_clock::now(); auto ms_comb = std::chrono::duration_cast(t1_comb - t0_comb).count(); + if (Ret != cudaSuccess) { + std::fprintf( + stderr, + "%s:%d: Failed to launch flash_fwd_splitkv_combine_kernel_jit: %s\n", + __FILE__, __LINE__, + static_cast(static_cast(Ret))); + std::abort(); + } std::fprintf(stderr, "[proteus-jit] %s h=%d combine_launch %ld ms\n", DtypeLabel, Config->HeadDim, ms_comb); C10_CUDA_KERNEL_LAUNCH_CHECK(); diff --git a/setup.py b/setup.py index c8054fc8ad2..1182a237fc6 100644 --- a/setup.py +++ b/setup.py @@ -313,6 +313,7 @@ def validate_and_update_archs(archs): "cuda", "cudart_static", "nvptxcompiler_static", + "nvrtc", "dl", "pthread", "rt",