[pat] (2/5) Add QKGrouper and refactor AttentionHeadGrouper for vmap

[lisjin]

Summary

• QKGrouper: new grouper that prunes only the Q and K slabs of a packed QKV tensor, leaving V untouched.
• AttentionHeadGrouperDim0/Dim1: stop in-place-mutating self.p.data and reassign self.p instead, so torch.vmap (used by PruneOptimizer's local_map path from PR1) can iterate over the grouped view without tripping the in-place check.

Changes

• torchao/prototype/pat/group/attention.py
• torchao/prototype/pat/group/__init__.py: re-export QKGrouper.

Depends on #4377

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/4378

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• Run pull request jobs on OSDC runners in shadow mode

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[lisjin]

@claude Please review changes

[claude]

Claude finished @lisjin's task in 1m 25s —— View job

---

Review of PR #4378

• Gather context and read changed files
• Review code changes
• Post review feedback

---

Overall: LGTM. The vmap refactor is clean and correct. One minor suggestion below.

---

AttentionHeadGrouperDim0 / Dim1 refactor

The change from self.p.data = self.p.data.view(...) to self.p = self.p.view(...) is the right approach for vmap compatibility. Since self._param retains the original tensor (set in Grouper.__init__), the base __exit__ correctly restores state via self.p = self._param without needing the explicit reshape-back that the old Dim0.__exit__ performed.

For Dim1, the custom __exit__ is still needed for the transpose+copy-back to self._orig_p, and the refactored version correctly uses a local data variable instead of mutating self.p.data — clean.

QKGrouper

The MRO and super() calls look correct:

1. PackedGrouperMixin.__init__ sets embed_dim / pack_dim
2. Dim1Grouper.__init__ (via super(PackedGrouperMixin, self)) wires up the slice through Grouper.__init__

One minor observation:

qk_reg_index validation (attention.py:32-35): The else branch silently treats any non-zero value as "key". Consider either an explicit elif qk_reg_index == 1 with a ValueError for invalid values, or just an assert:

assert qk_reg_index in (0, 1), f"qk_reg_index must be 0 (query) or 1 (key), got {qk_reg_index}"

This would prevent confusing behavior if someone mistakenly passes 2 expecting the value slice. Very minor — up to you whether it's worth adding for a prototype API.

---

No other issues found. The changes are minimal, well-motivated, and correctly preserve the grouper semantics while eliminating the in-place .data mutations that conflict with vmap.