Prebuild_one_add #1760
Prebuild_one_add #1760
Conversation
There was a problem hiding this comment.
Pull request overview
This PR enhances the prebuild infrastructure by implementing incremental builds and adding specialized MHA (Multi-Head Attention) variant prebuilding for PREBUILD_KERNELS mode 1. The changes optimize build times by skipping already-built modules and enable prebuilding of specific MHA configurations that are commonly used at runtime.
Key changes:
- Implemented incremental build logic to skip already-built modules, reducing rebuild time
- Added specialized prebuilding of two MHA varlen_fwd variants with specific flag configurations
- Modified module loading logic to support preferred custom module names with fallback behavior
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 10 comments.
| File | Description |
|---|---|
| setup.py | Adds incremental build support, implements PREBUILD_KERNELS==1 specific MHA variant building, refines module exclusion logic, and comments out cleanup code |
| aiter/jit/core.py | Enhances module loading to try custom module names first before falling back to default module name |
Comments suppressed due to low confidence (1)
setup.py:229
- The commented-out code on line 229 should be removed rather than left in the codebase. Having both the old and new implementation side-by-side creates confusion and reduces code maintainability.
class NinjaBuildExtension(BuildExtension):
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| else: | ||
| all_opts_args_build, _ = core.get_args_of_build("all", exclude=exclude_ops) | ||
|
|
||
| bd = f"{core.get_user_jit_dir()}/build" |
There was a problem hiding this comment.
The variable 'bd' is defined but never used, especially after the cleanup code that would have used it was commented out. This creates unnecessary code clutter and should be removed.
aiter/jit/core.py
Outdated
| else md_name | ||
| ) | ||
| module = get_module(preferred_md) | ||
| except Exception: |
There was a problem hiding this comment.
The exception variable 'e' on line 822 is captured but never used. This creates unnecessary code clutter. Either use the exception for logging/debugging purposes or remove the variable assignment and use a bare 'except Exception:'.
aiter/jit/core.py
Outdated
| else md_name | ||
| ) | ||
| module = get_module(preferred_md) | ||
| except Exception: |
There was a problem hiding this comment.
The module loading logic has been restructured but the error handling behavior has changed subtly. The new logic tries to load 'preferred_md' first (custom module name if gen_func exists), then falls back to 'md_name' if that fails. However, this means if a custom module name is specified but doesn't exist, it will silently fall back to the original module without any indication that the custom module was not found. Consider logging a warning when falling back to md_name after failing to load preferred_md to make debugging easier.
| except Exception: | |
| except Exception as exc: | |
| if preferred_md != md_name and isinstance(exc, ModuleNotFoundError): | |
| logging.getLogger(__name__).warning( | |
| "Preferred module '%s' not found, falling back to '%s'.", | |
| preferred_md, | |
| md_name, | |
| ) |
setup.py
Outdated
| f'{ck_dir}/example/ck_tile/01_fmha/generate.py -d fwd_splitkv --receipt 200 --filter " @ " --output_dir {{}}', | ||
| # AITER combined API generation for varlen fwd | ||
| f"{core.AITER_CSRC_DIR}/cpp_itfs/mha_fwd_generate.py --receipt 3 --output_dir {{}}", | ||
| ] | ||
| all_opts_args_build.append(variant_args) | ||
| # Also prebuild splitkv path to mirror JIT fallback behavior | ||
| variant_args_splitkv = dict(base_args) | ||
| variant_args_splitkv["md_name"] = ( | ||
| "mha_varlen_fwd_bf16_nlogits_nbias_nmask_lse_ndropout_nskip_nqscale" | ||
| ) | ||
| variant_args_splitkv["blob_gen_cmd"] = [ | ||
| f"{ck_dir}/example/ck_tile/01_fmha/generate.py -d fwd --receipt 200 --filter *bf16*_nlogits*_nbias*_nmask*_lse*_ndropout*_nskip*_nqscale* --output_dir {{}}", | ||
| f'{ck_dir}/example/ck_tile/01_fmha/generate.py -d fwd_splitkv --receipt 200 --filter " @ " --output_dir {{}}', |
There was a problem hiding this comment.
The filter pattern ' @ ' (with spaces) on lines 175 and 187 appears unusual and potentially incorrect. This pattern seems like it would match very few (if any) valid filenames. Verify that this is the intended filter pattern or if it should be a different pattern such as '@' or if the spaces are intentional for some specific matching behavior.
Motivation
Technical Details
Test Plan
Test Result
Submission Checklist