Conversation
|
Caution Review failedPull request was closed or merged during review 📝 WalkthroughWalkthroughThis PR introduces AOT build profile support enabling selection between full and lightweight ("edge_fm") kernel generation, adds CUDA graph compatibility via dynamic KV lengths, extends GQA group size support to 6 and 7, and updates build scripts to auto-detect profiles and CUDA architectures. ChangesAOT Build Profile System and CUDA Graph Enhancements
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 3❌ Failed checks (3 warnings)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces AOT build profiles to customize kernel compilation and adds CUDA graph support for single decode kernels by allowing dynamic sequence lengths via device pointers. It also optimizes GQA performance for specific group sizes on sm80 and removes a Boost dependency in the FP16 headers. Review feedback suggests centralizing the CI-detection logic for build profiles to reduce duplication, optimizing configuration management by using shallow copies instead of deep copies, and improving code maintainability by replacing magic numbers with named constants.
| - activation / norm / rope / page / sampling / topk helpers | ||
|
|
||
| It intentionally excludes heavyweight optional families such as XQA, MoE, | ||
| communication kernels, FP8 attention variants, and large head-dim matrices | ||
| that are not used in edge-fm today. | ||
| """ | ||
| return { | ||
| "fa2_head_dim": [(64, 64), (128, 128)], | ||
| "fa3_head_dim": [(64, 64), (128, 128)], | ||
| "f16_dtype": [torch.float16, torch.bfloat16], | ||
| "f8_dtype": [], |
There was a problem hiding this comment.
The logic for determining the default build profile (checking for CI environments to decide between "full" and "edge_fm") is currently duplicated in build_backend.py and several shell scripts. It would be more maintainable to centralize this logic here within normalize_build_profile.
| - activation / norm / rope / page / sampling / topk helpers | |
| It intentionally excludes heavyweight optional families such as XQA, MoE, | |
| communication kernels, FP8 attention variants, and large head-dim matrices | |
| that are not used in edge-fm today. | |
| """ | |
| return { | |
| "fa2_head_dim": [(64, 64), (128, 128)], | |
| "fa3_head_dim": [(64, 64), (128, 128)], | |
| "f16_dtype": [torch.float16, torch.bfloat16], | |
| "f8_dtype": [], | |
| def normalize_build_profile(profile: Optional[str]) -> str: | |
| profile_name = (profile or os.environ.get("FLASHINFER_AOT_BUILD_PROFILE")) | |
| if profile_name is None: | |
| if any( | |
| os.environ.get(var) | |
| for var in ("CI", "GITHUB_ACTIONS", "JENKINS_HOME", "JENKINS_URL") | |
| ): | |
| profile_name = "full" | |
| else: | |
| profile_name = "edge_fm" | |
| profile_name = profile_name.strip().lower().replace("-", "_") | |
| aliases = { | |
| "default": "full", | |
| "minimal": "edge_fm", | |
| "fast": "edge_fm", | |
| "dev": "edge_fm", | |
| "edgefm": "edge_fm", | |
| } | |
| return aliases.get(profile_name, profile_name) |
| "default": "full", | ||
| "minimal": "edge_fm", | ||
| "fast": "edge_fm", | ||
| "dev": "edge_fm", |
There was a problem hiding this comment.
copy.deepcopy is likely unnecessary here because get_config_for_profile returns a fresh dictionary created by get_default_config or get_edge_fm_fast_config. A simple assignment or .copy() would be more efficient.
| "dev": "edge_fm", | |
| final_config = get_config_for_profile(normalized).copy() |
| if any( | ||
| os.environ.get(var) | ||
| for var in ("CI", "GITHUB_ACTIONS", "JENKINS_HOME", "JENKINS_URL") | ||
| ): | ||
| build_profile = "full" | ||
| else: | ||
| build_profile = "edge_fm" | ||
| os.environ["FLASHINFER_AOT_BUILD_PROFILE"] = build_profile | ||
|
|
||
| # Set up build directory |
There was a problem hiding this comment.
This logic for determining the default build profile is duplicated. If normalize_build_profile in flashinfer/aot.py is updated to handle CI-aware defaults, this block can be simplified to rely on that centralized logic.
build_profile = aot.normalize_build_profile(None)
os.environ["FLASHINFER_AOT_BUILD_PROFILE"] = build_profile| o_ptr[i] = DTypeO(0); | ||
| } | ||
| if (lse != nullptr && tx == 0) { | ||
| lse[kv_chunk_idx * num_qo_heads + qo_head_idx] = -5e4; |
There was a problem hiding this comment.
The magic number -5e4 is used here as a very small LSE value to ensure this chunk has no contribution during state merging. It would be clearer to use a named constant or at least add a comment explaining why this specific value was chosen (e.g., to avoid NaN issues with -inf during reduction in MergeStates).
2 participants
📌 Description
🔍 Related Issues
🚀 Pull Request Checklist
Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.
✅ Pre-commit Checks
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
unittest, etc.).Reviewer Notes
Summary by CodeRabbit
Chores
Performance