Update: expand auto_incore scope to cover RoPE and cache update in Qwen3 decode#16
Conversation
…en3 decode - Lift `pl.auto_incore()` to wrap the entire Scope 2 (RoPE, KV cache update, and decode attention) instead of attention only - Enable save_kernels in compile_and_run for kernel inspection - Reformat multi-arg slice/assemble calls for readability
📝 WalkthroughWalkthroughWraps per-batch processing loops with Changes
Estimated code review effort🎯 2 (Simple) | ⏱️ ~10 minutes Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. 📝 Coding Plan
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 Tip CodeRabbit can scan for known vulnerabilities in your dependencies using OSV Scanner.OSV Scanner will automatically detect and report security vulnerabilities in your project's dependencies. No additional configuration is required. |
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request primarily focuses on enhancing the performance and debuggability of the Qwen3 decode model. It broadens the application of Highlights
Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![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 expands the auto_incore scope in qwen3_32b_decode.py to include RoPE and KV cache updates, which is a good optimization. It also improves code formatting and enables kernel saving for inspection. My feedback focuses on making the kernel saving feature more configurable to avoid surprising users with generated files.
| device_id=device_id, | ||
| rtol=2e-2, | ||
| atol=2e-2, | ||
| save_kernels=True, |
There was a problem hiding this comment.
Hardcoding save_kernels=True will cause this example to always generate kernel files, which might be an unintended side effect for users. It would be better to make this behavior configurable. For instance, you could add a save_kernels parameter to the compile_and_run function and control it from the calling context (e.g., via a command-line argument in the __main__ block).
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
examples/qwen3_32b_decode.py (1)
458-458: Make kernel dumping configurable instead of always enabled.Line 458 hard-enables kernel artifact dumping for all runs. This is great for debugging, but noisy for normal execution. Prefer a function parameter with a safer default.
♻️ Proposed diff
def compile_and_run( @@ dump_passes: bool = True, + save_kernels: bool = False, ): @@ config=RunConfig( @@ - save_kernels=True, + save_kernels=save_kernels, strategy=OptimizationStrategy.Default,🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@examples/qwen3_32b_decode.py` at line 458, The call currently hard-sets save_kernels=True which always dumps kernel artifacts; change this to be controlled by a function parameter (e.g., add a parameter save_kernels: bool = False to the surrounding function like main/decode_runner) and pass that parameter through where save_kernels is used (replace the literal True with the parameter). Ensure the new parameter defaults to False so kernel dumping is opt-in and update any callers or CLI wrapper to expose the option if desired.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@examples/qwen3_32b_decode.py`:
- Line 458: The call currently hard-sets save_kernels=True which always dumps
kernel artifacts; change this to be controlled by a function parameter (e.g.,
add a parameter save_kernels: bool = False to the surrounding function like
main/decode_runner) and pass that parameter through where save_kernels is used
(replace the literal True with the parameter). Ensure the new parameter defaults
to False so kernel dumping is opt-in and update any callers or CLI wrapper to
expose the option if desired.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 6a57654d-1beb-4d3d-8b75-9688931bc08a
📒 Files selected for processing (1)
examples/qwen3_32b_decode.py
2 participants
pl.auto_incore()to wrap the entire Scope 2 (RoPE, KV cache update, and decode attention) instead of attention onlySummary by CodeRabbit
Refactor
Chores