Skip to content

bugs fix prebuild=1/2 #1516

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
zufayu wants to merge 3 commits into
base: main
Choose a base branch
from
Open

bugs fix prebuild=1/2 #1516

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
63668d8
bugs fix
amd-zfyu Dec 1, 2025
e2c3f3f
more bug for batch mha
amd-zfyu Dec 1, 2025
4138884
skip some so to save build time
amd-zfyu Dec 3, 2025
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
4 changes: 2 additions & 2 deletions aiter/jit/optCompilerConfig.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -833,7 +833,7 @@
"'-DCK_TILE_FMHA_FWD_FAST_EXP2=1'",
"f'-DCK_TILE_FLOAT_TO_BFLOAT16_DEFAULT={os.environ.get(\"CK_TILE_FLOAT_TO_BFLOAT16_DEFAULT\", 2)}'",
"f'-DCK_TILE_ATTENTION_LOGITS_SOFT_CAP_DEFAULT={os.environ.get(\"CK_TILE_ATTENTION_LOGITS_SOFT_CAP_DEFAULT\", 0)}'",
"f'-DCK_TILE_ATTENTION_USE_SOFTSIGN_ASM={os.environ.get(\"CK_TILE_ATTENTION_USE_SOFTSIGN_ASM\", 1)}'"
"'-DCK_TILE_ATTENTION_USE_SOFTSIGN_ASM=0'"
],
"extra_ldflags": "None",
"extra_include": [
Expand Down Expand Up @@ -970,7 +970,7 @@
"blob_gen_cmd": [
"f'{CK_DIR}/example/ck_tile/01_fmha/generate.py -d fwd --receipt 600 --output_dir {{}}'",
"f'{CK_DIR}/example/ck_tile/01_fmha/generate.py -d fwd_splitkv --receipt 600 --output_dir {{}}'",
"f'{CK_DIR}/example/ck_tile/01_fmha/generate.py -d batch_prefill --receipt 600 --output_dir {{}}'",
"f'{CK_DIR}/example/ck_tile/01_fmha/generate.py -d batch_prefill --receipt 600 --optdim 128 --filter \"fmha_batch_prefill_d128_bf16_*_ndropout_*_nbias_*_nmask_*_nlse_*\" --output_dir {{}}'",
"f'{AITER_CSRC_DIR}/cpp_itfs/mha_fwd_generate.py --receipt 5 --output_dir {{}}'",
"f'{get_asm_dir()}/fmha_v3_fwd/codegen.py --output_dir {{}}'"
]
Expand Down
20 changes: 11 additions & 9 deletions setup.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
# !!!!!!!!!!!!!!!! never import aiter
# from aiter.jit import core
this_dir = os.path.dirname(os.path.abspath(__file__))
os.environ["AITER_META_DIR"] = this_dir
sys.path.insert(0, f"{this_dir}/aiter/")
from concurrent.futures import ThreadPoolExecutor

Expand Down Expand Up @@ -73,13 +74,6 @@ def is_develop_mode():
ck_dir
), 'CK is needed by aiter, please make sure clone by "git clone --recursive https://github.com/ROCm/aiter.git" or "git submodule sync ; git submodule update --init --recursive"'

if os.path.exists("aiter_meta") and os.path.isdir("aiter_meta"):
shutil.rmtree("aiter_meta")
shutil.copytree("3rdparty", "aiter_meta/3rdparty")
shutil.copytree("hsa", "aiter_meta/hsa")
shutil.copytree("gradlib", "aiter_meta/gradlib")
shutil.copytree("csrc", "aiter_meta/csrc")

def get_exclude_ops():
if PREBUILD_KERNELS == 1:
return [
Expand Down Expand Up @@ -174,6 +168,8 @@ def get_exclude_ops():
"module_mla_metadata",
"module_mla_reduce",
]
elif PREBUILD_KERNELS == 4:
return []
Copy link

Copilot AI Dec 1, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The get_exclude_ops() function does not return a value when PREBUILD_KERNELS is not 1, 2, or 3. This will cause exclude_ops to be None on line 104, leading to a potential error when passed to core.get_args_of_build().

Add a default return statement:

def get_exclude_ops():
    if PREBUILD_KERNELS == 1:
        return [...]
    elif PREBUILD_KERNELS == 2:
        return [...]
    elif PREBUILD_KERNELS == 3:
        return []
    else:
        return []
Suggested change
return []
return []
else:
return []

Copilot uses AI. Check for mistakes.
else:
return []

Expand Down Expand Up @@ -225,7 +221,13 @@ def build_one_module(one_opt_args):
raise NotImplementedError("Only ROCM is supported")


# aiter_meta prepared above
if os.path.exists("aiter_meta") and os.path.isdir("aiter_meta"):
shutil.rmtree("aiter_meta")
## link "3rdparty", "hsa", "csrc" into "aiter_meta"
shutil.copytree("3rdparty", "aiter_meta/3rdparty")
shutil.copytree("hsa", "aiter_meta/hsa")
shutil.copytree("gradlib", "aiter_meta/gradlib")
shutil.copytree("csrc", "aiter_meta/csrc")


class NinjaBuildExtension(BuildExtension):
Expand Down Expand Up @@ -293,4 +295,4 @@ def has_ext_modules(self):
)

if os.path.exists("aiter_meta") and os.path.isdir("aiter_meta"):
shutil.rmtree("aiter_meta")
shutil.rmtree("aiter_meta")
Loading