test: pin Llama3 TurboAttention GQA layout (PRPUNDIT-20)

[jiagaoxiang]

This PR was created by dougljia via Test Gap Resolver.

Summary

• Adds tests/unit_tests/backends/torchtitan/test_llama3_turbo_attention.py.
• Stub inner_attention as an nn.Module and assert unrepeated KV heads in BSHD (no repeat_kv / BHSD transpose), positions forwarded to RoPE, and that output still goes through wo.

Closes test gap PRPUNDIT-20.

Test plan

• PYTHONPATH=. pytest tests/unit_tests/backends/torchtitan/test_llama3_turbo_attention.py (5 passed, CPU, needs torchtitan)
• black --check / isort --profile black on the new file

[Copilot]

🟡 Changes recommended

One test hard-codes global head counts despite the implementation explicitly inferring local heads for tensor-parallel sharding, which can cause false failures in valid configurations.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

Adds unit coverage to pin the expected tensor layout and argument forwarding in the Primus LLaMA3 TurboAttention Attention.forward mirror, ensuring GQA KV heads are not expanded externally and that RoPE positions and wo projection routing remain intact.

Changes:

• Adds tests/unit_tests/backends/torchtitan/test_llama3_turbo_attention.py with stubs/spies to isolate and validate the mirror forward behavior.
• Verifies unrepeated KV head shapes (BSHD), positions forwarding into apply_rotary_emb, and that outputs still pass through Attention.wo.

File summaries

File	Description
tests/unit_tests/backends/torchtitan/test_llama3_turbo_attention.py	New unit tests to lock in Primus LLaMA3 TurboAttention mirror forward layout and API behavior.

Review details

• Files reviewed: 1/1 changed files
• Comments generated: 1
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[jiagaoxiang]

TestIntel PR Steward: Looked at the failed build-docker (3.12) check on head 54d8bc0.

The Docker image itself builds fine (all layers cached, image exported successfully). The job then fails at:

docker login -u tasimage -p 
flag needs an argument: 'p' in -p
##[error]Process completed with exit code 125.

i.e. the registry password is empty. That's consistent with GitHub Actions withholding repo secrets from pull_request-triggered workflows on fork branches — this PR runs from jiagaoxiang/Primus (a fork). Checking sibling Test Gap Resolver PRs confirms the pattern: every testgap/* fork branch (PRPUNDIT-18, 19, 20, 21, 22, 23) hits the same build-docker failure, while non-fork branches don't show this symptom. This is a repo-wide CI/secrets configuration issue, not something introduced by this PR's test file, and nothing in tests/unit_tests/backends/torchtitan/test_llama3_turbo_attention.py can fix it. Rerunning won't help either, since fork PRs never get secret access regardless of attempt count.

Flagging for a maintainer to address at the workflow/secrets level (e.g. skip registry push for fork PRs, or use pull_request_target carefully). Not making any changes here and not merging while this check is red.

Also replied on the Copilot review thread re: the GQA head-count assertion.

[jiagaoxiang]

Replying to Copilot's finding on the GQA head-count assertion.

[jiagaoxiang]

TestIntel PR Steward: Thanks — I looked at whether this test can actually see the failure mode you describe.

This test instantiates Attention(llama3_args) directly with no tensor-parallel mesh or sharding applied anywhere in the test (no DeviceMesh, no parallelize_module, no distributed process group). wq/wk/wv are therefore always built with the global n_heads/n_kv_heads from llama3_args, so xq.view(bs, seqlen, -1, head_dim) always resolves to llama3_args.n_heads here — the -1 and the global arg are the same number in this unit test's scope. So the assertions against llama3_args.n_heads/n_kv_heads are correct for what this test actually exercises.

Deriving the expected counts from wq(x)/wk(x)/wv(x) shapes instead would be a defensive nice-to-have (guards against a future test that does add TP sharding), but it doesn't change correctness today, and this PR's scope is pinning the existing GQA-layout behavior for the test gap (PRPUNDIT-20), not adding TP coverage. Leaving the assertions as-is; happy to revisit if a follow-up test adds actual sharded execution.

[jiagaoxiang]

Superseded by same-repo PR #1103 (head AMD-AGI:testgap/PRPUNDIT-20) so Primus-CI-TAS receives Docker Hub credentials and can run torch unit tests. Closing this fork PR.