Update: unify paged attention golden cases across three test variants

[hw-native-sys-bot]

Summary

• paged_attention: Case2 batch 256→64, removed Case4/5/6
• paged_attention_unroll: Case2 batch 256→64
• batch_paged_attention: already matched, no changes needed

All three PA test variants (batch_paged_attention, paged_attention, paged_attention_unroll) now share identical Case1/2/3 configurations:

Case	QHEADS	KVHEADS	HEADDIM	BLOCKSIZE	context_len	batch
Case1	16	1	128	128	8192	256
Case2	64	1	128	64	8192	64
Case3	64	1	256	64	8192	64

Test plan

• Run batch_paged_attention device test with all 3 cases
• Run paged_attention device test with all 3 cases
• Run paged_attention_unroll device test with all 3 cases

[gemini-code-assist]

Summary of Changes

Hello, I'm Gemini Code Assist¹! 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 standardizes the golden test configurations across different paged attention implementations. By unifying the test cases for paged_attention, paged_attention_unroll, and batch_paged_attention, it aims to ensure consistency and simplify the testing and validation process for these components.

Highlights

• Paged Attention Test Cases: Modified the Case2 batch size from 256 to 64 and removed Case4, Case5, and Case6 in the paged_attention golden test configuration.
• Paged Attention Unroll Test Cases: Adjusted the Case2 batch size from 256 to 64 in the paged_attention_unroll golden test configuration.
• Test Case Unification: Ensured all three paged attention test variants (batch_paged_attention, paged_attention, paged_attention_unroll) now share identical Case1, Case2, and Case3 configurations for consistency.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The 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 /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

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 .gemini/ folder in the base of the repository. Detailed instructions can be found here.

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.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request unifies the golden test cases for paged_attention and paged_attention_unroll to match batch_paged_attention. The changes involve updating the batch size for Case2 and removing obsolete test cases (Case4, Case5, Case6) in paged_attention/golden.py, and updating the batch size for Case2 in paged_attention_unroll/golden.py. These modifications successfully align the test configurations as described, improving consistency across the test suite. The changes are straightforward and appear correct.