Generate kernel-aware NAS candidates#110
Conversation
|
Warning Review limit reached
Your plan currently allows 1 review/hour. Refill in 56 minutes and 30 seconds. Your organization has run out of usage credits. Purchase more in the billing tab. ⌛ How to resolve this issue?After more review capacity refills, a review can be triggered using the We recommend that you space out your commits to avoid hitting the rate limit. 🚦 How do rate limits work?CodeRabbit enforces hourly rate limits for each developer per organization. Our paid plans have higher rate limits than trial, open-source, and free plans. In all cases, review capacity refills continuously over time. Please see our FAQ for further information. ℹ️ Review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (4)
📝 WalkthroughWalkthroughThis PR introduces kernel-aware NAS candidate generation, adding a ChangesKernel-aware NAS candidate generation and ranking
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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 |
Darkroom4364
force-pushed
the
feat/nas-candidate-generation
branch
from
fb4c1fb to
d783032
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/research_engine/arch_cost_model.py`:
- Around line 131-145: The current use of truthy fallbacks (e.g., hidden_dims =
hidden_dims or [base_config["hidden_dim"]]) hides explicit empty-sweep lists and
the tile_multiple parameter can cause a ZeroDivisionError when <= 0; change each
fallback to use "is None" checks so empty lists are preserved (e.g., if
hidden_dims is None: hidden_dims = [...]) for hidden_dims, head_dims,
ffn_ratios, kv_head_counts, window_sizes, and qk_norm_options, and add
validation for tile_multiple at the start (raise a ValueError if tile_multiple
is <= 0) so downstream code that divides by tile_multiple cannot error
unpredictably.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 06052804-503c-45b3-8161-aa7b46400b62
📒 Files selected for processing (4)
docs/research/kernel-aware-nas.mdscripts/nas_experiment.pysrc/research_engine/arch_cost_model.pytests/test_arch_cost_model.py
| tile_multiple: int = MATMUL_TILE_SIZE, | ||
| max_candidates: int | None = None, | ||
| ) -> list[dict[str, Any]]: | ||
| """Generate tile-aligned transformer architecture candidates. | ||
|
|
||
| The search varies dimensions that are natural handles for kernel-aware NAS: | ||
| hidden width, attention head width, GQA ratio, FFN expansion, sliding-window | ||
| attention, and QK-norm placement. Invalid combinations are skipped. | ||
| """ | ||
| hidden_dims = hidden_dims or [base_config["hidden_dim"]] | ||
| head_dims = head_dims or [64, 128, 256] | ||
| ffn_ratios = ffn_ratios or [3.0, 4.0, 5.333] | ||
| kv_head_counts = kv_head_counts or [1, 2, 4, 8] | ||
| window_sizes = window_sizes or [base_config.get("window_size"), None] | ||
| qk_norm_options = qk_norm_options or [base_config.get("use_qk_norm", True), False] |
There was a problem hiding this comment.
Handle explicit empty sweeps and invalid tile_multiple deterministically.
Line 140-145 currently uses truthy fallback (or), so explicit empty lists are ignored and silently replaced by defaults. Also, Line 160 can raise ZeroDivisionError when tile_multiple <= 0. Please validate the argument and switch to is None defaults.
Suggested patch
def generate_nas_candidates(
@@
) -> list[dict[str, Any]]:
@@
- hidden_dims = hidden_dims or [base_config["hidden_dim"]]
- head_dims = head_dims or [64, 128, 256]
- ffn_ratios = ffn_ratios or [3.0, 4.0, 5.333]
- kv_head_counts = kv_head_counts or [1, 2, 4, 8]
- window_sizes = window_sizes or [base_config.get("window_size"), None]
- qk_norm_options = qk_norm_options or [base_config.get("use_qk_norm", True), False]
+ if tile_multiple <= 0:
+ raise ValueError("tile_multiple must be > 0")
+
+ hidden_dims = [base_config["hidden_dim"]] if hidden_dims is None else hidden_dims
+ head_dims = [64, 128, 256] if head_dims is None else head_dims
+ ffn_ratios = [3.0, 4.0, 5.333] if ffn_ratios is None else ffn_ratios
+ kv_head_counts = [1, 2, 4, 8] if kv_head_counts is None else kv_head_counts
+ window_sizes = [base_config.get("window_size"), None] if window_sizes is None else window_sizes
+ qk_norm_options = [base_config.get("use_qk_norm", True), False] if qk_norm_options is None else qk_norm_optionsAlso applies to: 160-160
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/research_engine/arch_cost_model.py` around lines 131 - 145, The current
use of truthy fallbacks (e.g., hidden_dims = hidden_dims or
[base_config["hidden_dim"]]) hides explicit empty-sweep lists and the
tile_multiple parameter can cause a ZeroDivisionError when <= 0; change each
fallback to use "is None" checks so empty lists are preserved (e.g., if
hidden_dims is None: hidden_dims = [...]) for hidden_dims, head_dims,
ffn_ratios, kv_head_counts, window_sizes, and qk_norm_options, and add
validation for tile_multiple at the start (raise a ValueError if tile_multiple
is <= 0) so downstream code that divides by tile_multiple cannot error
unpredictably.
2 participants
Summary
Refs #81
Tests
Note: ruff is not installed locally.
Summary by CodeRabbit
New Features
Documentation
Tests