Skip to content

Update for proteus PR https://github.com/Olympus-HPC/proteus/pull/429 #1

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ggeorgakoudis wants to merge 2 commits into
base: proteus
Choose a base branch
from
Open

Update for proteus PR https://github.com/Olympus-HPC/proteus/pull/429 #1

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
663dfb6
Update for proteus PR https://github.com/Olympus-HPC/proteus/pull/429
ggeorgakoudis Mar 13, 2026
711b170
Updates
ggeorgakoudis Mar 20, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 20 additions & 14 deletions benchmarks/benchmark_flash_attention.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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

Expand Down
16 changes: 10 additions & 6 deletions build_fa_proteus.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)}"
Expand Down Expand Up @@ -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}"
Expand All @@ -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}"
Expand All @@ -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}"
Expand All @@ -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)
Expand Down
60 changes: 45 additions & 15 deletions csrc/flash_attn/src/flash_bwd_jit_bridge.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -218,7 +218,8 @@ bool use_proteus_jit() {
return Enabled;
}

std::vector<std::string> get_extra_args() {
std::vector<std::string> 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);
Expand All @@ -231,7 +232,16 @@ std::vector<std::string> 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) {
Expand Down Expand Up @@ -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<proteus::CppJitModule>(
"cuda", JitKernelCode, get_extra_args());
"cuda", JitKernelCode, get_extra_args(Backend), Backend);
return *Module;
}

Expand Down Expand Up @@ -540,12 +551,18 @@ bool try_run_mha_bwd_jit_impl(Flash_bwd_params &params, 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);
Instance.launch(GridM, Block, /*smem=*/0,
reinterpret_cast<void *>(stream), params);
auto Ret = Instance.launch(GridM, Block, /*smem=*/0,
reinterpret_cast<void *>(stream), params);
if (Ret != cudaSuccess) {
std::fprintf(stderr, "%s:%d: Kernel 1 launch failed with error code %s\n",
__FILE__, __LINE__,
cudaGetErrorString(static_cast<cudaError_t>(static_cast<int>(Ret))));
std::abort();
}
C10_CUDA_KERNEL_LAUNCH_CHECK();
}

Expand All @@ -554,14 +571,22 @@ bool try_run_mha_bwd_jit_impl(Flash_bwd_params &params, 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),
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<void *>(stream), params);
Instance.setFuncAttribute(
proteus::CppJitFuncAttribute::MaxDynamicSharedMemorySize, SmemSize);
auto Ret = Instance.launch(GridN, Block, SmemSize,
reinterpret_cast<void *>(stream), params);
if (Ret != cudaSuccess) {
std::fprintf(stderr, "%s:%d: Kernel 1 launch failed with error code %s\n",
__FILE__, __LINE__,
cudaGetErrorString(static_cast<cudaError_t>(static_cast<int>(Ret))));
std::abort();
}
C10_CUDA_KERNEL_LAUNCH_CHECK();
}

Expand All @@ -571,10 +596,19 @@ bool try_run_mha_bwd_jit_impl(Flash_bwd_params &params, 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<void *>(stream), params, nsplits);
Instance.setFuncAttribute(
proteus::CppJitFuncAttribute::MaxDynamicSharedMemorySize, SmemSize);
auto Ret =
Instance.launch(GridM, Block, SmemSize,
reinterpret_cast<void *>(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<cudaError_t>(static_cast<int>(Ret))));
std::abort();
}
C10_CUDA_KERNEL_LAUNCH_CHECK();
}

Expand All @@ -596,8 +630,6 @@ bool try_run_mha_bwd_jit_fp16(Flash_bwd_params &params, 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<std::chrono::milliseconds>(t1 - t0).count();
std::fprintf(stderr, "[proteus-jit] bwd_fp16 h=%d launch %ld ms\n",
Config ? Config->HeadDim : params.d, ms);
return result;
}

Expand All @@ -610,8 +642,6 @@ bool try_run_mha_bwd_jit_bf16(Flash_bwd_params &params, 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<std::chrono::milliseconds>(t1 - t0).count();
std::fprintf(stderr, "[proteus-jit] bwd_bf16 h=%d launch %ld ms\n",
Config ? Config->HeadDim : params.d, ms);
return result;
}

Expand Down
45 changes: 34 additions & 11 deletions csrc/flash_attn/src/flash_fwd_jit_bridge.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -236,7 +236,8 @@ bool use_proteus_jit() {
return Enabled;
}

std::vector<std::string> get_extra_args() {
std::vector<std::string> 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);
Expand All @@ -249,7 +250,16 @@ std::vector<std::string> 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__",
Comment thread
ZwFink marked this conversation as resolved.
};
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) {
Expand All @@ -261,8 +271,9 @@ std::vector<std::string> 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<proteus::CppJitModule>(
"cuda", JitKernelCode, get_extra_args());
"cuda", JitKernelCode, get_extra_args(Backend), Backend);
return *Module;
}

Expand Down Expand Up @@ -454,20 +465,27 @@ bool try_run_mha_fwd_jit_fp16(Flash_fwd_params &params, 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),
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<void *>(stream),
params);
auto Ret = Instance.launch(Grid, Block, SmemSize,
reinterpret_cast<void *>(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<cudaError_t>(static_cast<int>(Ret))));
std::abort();
}

auto ms = std::chrono::duration_cast<std::chrono::milliseconds>(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;
}
Expand Down Expand Up @@ -510,20 +528,25 @@ bool try_run_mha_fwd_jit_bf16(Flash_fwd_params &params, 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),
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<void *>(stream),
auto Ret = Instance.launch(Grid, Block, SmemSize, reinterpret_cast<void *>(stream),
params);
auto t1 = std::chrono::steady_clock::now();
if (static_cast<int>(Ret) != cudaSuccess) {
std::fprintf(stderr, "%s:%d: Failed to launch flash_fwd_kernel_jit: %s\n",
__FILE__, __LINE__,
cudaGetErrorString(static_cast<cudaError_t>(static_cast<int>(Ret))));
std::abort();
}
auto ms = std::chrono::duration_cast<std::chrono::milliseconds>(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;
}
Expand Down
Loading