Add required keys to attention pruning config

[grzegorz-k-karch]

What does this PR do?

Type of change: ? Bug fix

The config examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/pruning/attn_pruning.yaml didn't have required keys to use attention pruning in the example examples/puzzletron/main.py

Usage

Testing

In examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/Llama-3_1-8B.yaml
change ffn_pruning to attn_pruning

Before your PR is "Ready for review"

Make sure you read and follow Contributor guidelines and your commits are signed (git commit -s -S).

Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded trust_remote_code=True, torch.load(..., weights_only=False), pickle, etc.).

• Is this change backward compatible?: ✅
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A
• Did you write any new necessary tests?: N/A
• Did you update Changelog?: N/A

Additional Information

Summary by CodeRabbit

• Chores
  • Updated pruning configuration for improved KV-head pruning support, including enhanced importance hook settings and attention output handling for memory optimization.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: ce0ba0c5-75ed-424e-9b2a-15abc7a16531

📥 Commits

Reviewing files that changed from the base of the PR and between 6e08b13 and d8c2f8c.

📒 Files selected for processing (1)

• examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/pruning/attn_pruning.yaml

---

📝 Walkthrough

Walkthrough

A configuration file for LLaMA pruning is updated to explicitly specify the importance hook class, add KV-head pruning settings via a pruning mixin, define the LLaMA KV-head layer descriptor, and designate the attention output projection target layer for pruning operations.

Changes

Cohort / File(s)	Summary
Pruning Configuration examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/pruning/attn_pruning.yaml	Added explicit importance hook class (IndependentKvHeadContributionHook), KV-head pruning mixin configuration (KVHeadsPruningMixIn), LLaMA KV-head layer descriptor, and attention output projection target layer specification (self_attn.o_proj).

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

🚥 Pre-merge checks | ✅ 6

✅ Passed checks (6 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title 'Add required keys to attention pruning config' directly and specifically describes the main change: adding required keys to the attention pruning configuration file.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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.
Security Anti-Patterns	✅ Passed	PR only modifies YAML configuration file with no Python code changes; no security anti-patterns detected.

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

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch gkarch/fix-attn-pruning-config

---

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

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 77.48%. Comparing base (6e08b13) to head (d8c2f8c).
⚠️ Report is 1 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1360      +/-   ##
==========================================
+ Coverage   76.93%   77.48%   +0.54%     
==========================================
  Files         471      471              
  Lines       50401    50401              
==========================================
+ Hits        38777    39052     +275     
+ Misses      11624    11349     -275     

Flag	Coverage Δ
examples	41.59% <ø> (+0.92%)	⬆️
unit	52.74% <ø> (ø)

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[kevalmorabia97]

/ok to test d8c2f8c

[github-actions]

PR Preview Action v1.8.1
Preview removed because the pull request was closed.
2026-04-28 13:05 UTC