From 03226ac5b434b603236224e48b5682099bd56fbb Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Wed, 22 Jul 2026 08:52:49 -0700 Subject: [PATCH 1/3] [None][fix] Do not treat a bare fmha_v2_cu directory as completed generation generate_fmha_cu creates the fmha_v2_cu output directory before running the (minutes-long) fmha_v2 kernel generation, but the build skips generation whenever that directory exists. An interrupted first generation (Ctrl-C, node failure, build timeout) therefore leaves an empty or partial directory that every subsequent non-clean build trusts: the launcher .cu sources are missing while the checked-in fmha_cubin.cpp table still references them, and the build fails much later at link with undefined run_fmha_v2_* symbols. Nothing recovers short of --clean, which is not an obvious remedy for that error. Write a .generation_complete stamp as the last step of generate_fmha_cu (removed again as its first step) and key the skip check on the stamp instead of the directory. Existing checkouts regenerate once to acquire the stamp; the content-compare copy keeps mtimes of unchanged files, so no recompilation results. Signed-off-by: Brian Nguyen --- scripts/build_wheel.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/scripts/build_wheel.py b/scripts/build_wheel.py index 1ef38a753ea8..a8ba018f426d 100755 --- a/scripts/build_wheel.py +++ b/scripts/build_wheel.py @@ -264,9 +264,18 @@ def setup_conan(scripts_dir, venv_python): return venv_conan +def fmha_generation_stamp(fmha_v2_cu_dir): + # 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" @@ -323,6 +332,8 @@ def move_if_updated(src, dst): if file_path not in generated_files: os.remove(file_path) + fmha_generation_stamp(fmha_v2_cu_dir).touch() + def create_cuda_stub_links(cuda_stub_dir: str, missing_libs: list[str]) -> str: """ @@ -678,7 +689,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): From 7d048a4dba21f9f6eafab08f5c59fef9300782cd Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Wed, 22 Jul 2026 11:23:03 -0500 Subject: [PATCH 2/3] [None][fix] Rename fmha_generation_stamp to private helper with type annotations Rename fmha_generation_stamp to _fmha_generation_stamp and annotate its parameter and return type with Path, per coding guidelines for private helpers and annotation coverage. Signed-off-by: Brian Nguyen --- scripts/build_wheel.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/scripts/build_wheel.py b/scripts/build_wheel.py index a8ba018f426d..b7912779370f 100755 --- a/scripts/build_wheel.py +++ b/scripts/build_wheel.py @@ -264,7 +264,7 @@ def setup_conan(scripts_dir, venv_python): return venv_conan -def fmha_generation_stamp(fmha_v2_cu_dir): +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 @@ -275,7 +275,7 @@ def fmha_generation_stamp(fmha_v2_cu_dir): 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_generation_stamp(fmha_v2_cu_dir).unlink(missing_ok=True) fmha_v2_dir = project_dir / "cpp/kernels/fmha_v2" @@ -332,7 +332,7 @@ def move_if_updated(src, dst): if file_path not in generated_files: os.remove(file_path) - fmha_generation_stamp(fmha_v2_cu_dir).touch() + _fmha_generation_stamp(fmha_v2_cu_dir).touch() def create_cuda_stub_links(cuda_stub_dir: str, missing_libs: list[str]) -> str: @@ -690,7 +690,7 @@ def main(*, fmha_v2_cu_dir = project_dir / "cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmha_v2_cu" if (clean or generate_fmha - or not fmha_generation_stamp(fmha_v2_cu_dir).exists()): + or not _fmha_generation_stamp(fmha_v2_cu_dir).exists()): generate_fmha_cu(project_dir, venv_python) with working_directory(build_dir): From b392900934bc1e24ca753ad15a1bd33ec8ebf79a Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Thu, 23 Jul 2026 14:46:48 -0500 Subject: [PATCH 3/3] [None][fix] Validate generated_files before writing FMHA completion stamp Raise RuntimeError if no *_sm*.cu files were produced before touching the stamp, so a silent generation failure doesn't leave an apparently complete stamp that causes subsequent builds to skip regeneration. Signed-off-by: Brian Nguyen --- scripts/build_wheel.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/scripts/build_wheel.py b/scripts/build_wheel.py index b7912779370f..462b74d8e511 100755 --- a/scripts/build_wheel.py +++ b/scripts/build_wheel.py @@ -325,6 +325,11 @@ 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: