Skip to content

[None][fix] Do not treat a bare fmha_v2_cu directory as completed generation #16743

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
brnguyen2 wants to merge 3 commits into
base: main
Choose a base branch
from
+18 −1
Open

[None][fix] Do not treat a bare fmha_v2_cu directory as completed generation #16743

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
03226ac
[None][fix] Do not treat a bare fmha_v2_cu directory as completed gen…
brnguyen2 Jul 22, 2026
7d048a4
[None][fix] Rename fmha_generation_stamp to private helper with type …
brnguyen2 Jul 22, 2026
b392900
[None][fix] Validate generated_files before writing FMHA completion s…
brnguyen2 Jul 23, 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
19 changes: 18 additions & 1 deletion scripts/build_wheel.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -264,9 +264,18 @@ def setup_conan(scripts_dir, venv_python):
return venv_conan


def _fmha_generation_stamp(fmha_v2_cu_dir: Path) -> Path:
# Written as the last step of generate_fmha_cu; its absence means a
# previous generation was interrupted and the directory contents cannot
# be trusted (the bare directory exists from the moment generation
# starts).
return fmha_v2_cu_dir / ".generation_complete"


def generate_fmha_cu(project_dir, venv_python):
fmha_v2_cu_dir = project_dir / "cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmha_v2_cu"
fmha_v2_cu_dir.mkdir(parents=True, exist_ok=True)
_fmha_generation_stamp(fmha_v2_cu_dir).unlink(missing_ok=True)

fmha_v2_dir = project_dir / "cpp/kernels/fmha_v2"

Expand Down Expand Up @@ -316,13 +325,20 @@ def move_if_updated(src, dst):
move_if_updated(cu_file, dst_file)
generated_files.add(str(dst_file.resolve()))

if not generated_files:
raise RuntimeError(
f"FMHA generation produced no *_sm*.cu files in {fmha_v2_cu_dir}; "
"generation may have failed silently.")

# Remove extra files
for root, _, files in os.walk(fmha_v2_cu_dir):
for file in files:
file_path = os.path.realpath(os.path.join(root, file))
if file_path not in generated_files:
os.remove(file_path)

_fmha_generation_stamp(fmha_v2_cu_dir).touch()

Comment thread
coderabbitai[bot] marked this conversation as resolved.

def create_cuda_stub_links(cuda_stub_dir: str, missing_libs: list[str]) -> str:
"""
Expand Down Expand Up @@ -678,7 +694,8 @@ def main(*,
source_dir = get_source_dir()

fmha_v2_cu_dir = project_dir / "cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmha_v2_cu"
if clean or generate_fmha or not fmha_v2_cu_dir.exists():
if (clean or generate_fmha
or not _fmha_generation_stamp(fmha_v2_cu_dir).exists()):
generate_fmha_cu(project_dir, venv_python)

with working_directory(build_dir):
Expand Down
Loading