Skip to content

refactor: migrate rms_norm + dsv4 attention to auto_chunk#386

Open
lyfne123 wants to merge 3 commits into
hw-native-sys:mainfrom
lyfne123:refactor/attention-rmsnorm-auto-chunk
Open

refactor: migrate rms_norm + dsv4 attention to auto_chunk#386
lyfne123 wants to merge 3 commits into
hw-native-sys:mainfrom
lyfne123:refactor/attention-rmsnorm-auto-chunk

Conversation

@lyfne123
Copy link
Copy Markdown
Contributor

Finishes #373 migration: rms_norm and dsv4 decode_attention_hca/swa were kept on chunked_loop_optimizer during #373 due to suspected sim mismatch; root cause was a known sim infra issue, unrelated to auto_chunk. Migrate the remaining 5 sites so the repo is fully on pl.auto_chunk.

🤖 Generated with Claude Code

sim failures were a known infra issue, unrelated to auto_chunk.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
@coderabbitai
Copy link
Copy Markdown

coderabbitai Bot commented May 26, 2026

Review Change Stack

Warning

Review limit reached

@lyfne123, we couldn't start this review because you've reached your PR review rate limit.

More reviews will be available in 26 minutes and 34 seconds. Learn how PR review limits work.

Your organization has run out of usage credits. Purchase more in the billing tab.

⌛ How to resolve this issue?

After more reviews become available, a review can be triggered using the @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

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 include higher PR review limits than trial, open-source, and free plans. In all cases, reviews become available again over time. During sustained high-volume PR review activity, CodeRabbit may temporarily slow when the next review becomes available.

Please see our Fair Usage Limits Policy for further information.

ℹ️ Review info
⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 00c2685c-9eb0-4df4-996f-17f1f20b8d7a

📥 Commits

Reviewing files that changed from the base of the PR and between 066d4d6 and 823e439.

📒 Files selected for processing (2)
  • models/deepseek/v4/decode_attention_hca.py
  • models/deepseek/v4/decode_attention_swa.py
📝 Walkthrough

Walkthrough

Five pl.at() loop optimizer annotations across RMS normalization, HCA, and SWA attention models are migrated from the deprecated optimization=pl.chunked_loop_optimizer single-parameter form to the new optimizations=[pl.auto_chunk] multi-parameter form.

Changes

Loop Optimizer Parameter Migration

Layer / File(s) Summary
RMS norm core group optimizer update
examples/intermediate/rms_norm.py
RMS normalization example's core-group tiling loop is updated to use optimizations=[pl.auto_chunk] instead of optimization=pl.chunked_loop_optimizer.
DeepSeek attention decoder optimizer updates
models/deepseek/v4/decode_attention_hca.py, models/deepseek/v4/decode_attention_swa.py
Three sparse-attention loops in HCA top-k generation and SWA KV scatter, sparse top-k, and block-table initialization paths switch to the new optimizations=[pl.auto_chunk] parameter form.

Estimated code review effort

🎯 1 (Trivial) | ⏱️ ~5 minutes

Possibly related PRs

  • hw-native-sys/pypto-lib#373: Broader refactor migrating the same pl.at() optimizer callsites from deprecated single-parameter to new multi-parameter form.

Poem

🐰 Loop optimizers hop and dance,
From old form to new—a parameter prance!
chunked_loop fades, auto_chunk takes flight,
Five lines of grace make compilation right! ✨

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 25.00% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Title check ✅ Passed The title accurately describes the main change: migrating rms_norm and dsv4 attention components from chunked_loop_optimizer to auto_chunk across multiple files.
Description check ✅ Passed The description is directly related to the changeset, explaining the context of finishing migration #373 and why these components are being migrated despite previous concerns.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.


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.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request replaces the deprecated optimization=pl.chunked_loop_optimizer with optimizations=[pl.auto_chunk] across several model files and examples. The feedback points out that in several locations (decode_attention_hca.py and decode_attention_swa.py), the pl.auto_chunk optimization is redundant because there are no loops inside the with pl.at blocks, and suggests removing the optimizations argument in those cases.

Comment thread models/deepseek/v4/decode_attention_hca.py Outdated
Comment thread models/deepseek/v4/decode_attention_swa.py Outdated
Comment thread models/deepseek/v4/decode_attention_swa.py Outdated
lyfne123 and others added 2 commits May 26, 2026 15:47
These pl.at blocks have no inner loop (chunk loop is outside), so
auto_chunk is a no-op; using it caused ~35% sim mismatch. Bare pl.at
per gemini review.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
bare pl.at: hca worsens to 80% mismatch, swa errors (scatter has
chunk= inner loop needing auto_chunk). only chunked_loop works for
these. 386 net effect = rms_norm migration only.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant