Use with AI · Chapter Application

Audit PyTorch Attention Shapes, QKV and Masks

Trace multi-head attention dimensions and masking semantics through a real implementation to find silent QKV, head, transpose and mask errors.

Chapter Application PyTorch: Zero to Hero PyTorch attention and transformer code Advanced
Read the source chapter → View the learning path →

How to use this

  1. Open a repository-aware AI assistant.
  2. Give it access to the repository or files you want reviewed.
  3. Copy the prompt below and run it unchanged first.
  4. Use the evidence it finds to decide what to inspect or change next.
PromptCopy and run against your own project
You are auditing a PyTorch attention or transformer implementation for dimensional and masking defects.

Do not assume a runtime-successful attention block is correct. Attention can execute while attending to the wrong positions, mixing heads incorrectly or applying masks with unintended broadcasting.

Trace one complete attention call from input embeddings through output projection.

Work in this order:

1. Record the dimensional contract for:
   - input embeddings
   - Q, K and V projections
   - head split
   - attention scores
   - masks
   - softmax weights
   - weighted values
   - head concatenation
   - output projection

Use semantic names such as batch, heads, query_length, key_length and head_dim rather than raw integers only.

2. Verify invariants:
   - embed_dim == num_heads * head_dim where required
   - Q/K inner dimensions align for score computation
   - attention scores have the intended [batch, heads, query, key] semantics
   - V multiplication returns the intended per-head representation
   - concatenation reconstructs the embedding dimension correctly

3. Inspect every reshape/view/transpose/permute and identify whether its semantic layout matches the next operation.

4. Audit masks separately:
   - causal mask orientation
   - padding mask meaning
   - boolean versus additive masks
   - mask rank and broadcasting
   - query length versus key length
   - cross-attention differences
   - masked values before softmax
   - all-masked rows producing NaNs

5. Compare custom attention code with any higher-level PyTorch API usage and flag incompatible assumptions about batch_first, mask shapes or return layouts.

Classify findings as:
- confirmed dimensional defect
- confirmed masking defect
- silent semantic risk
- fragile assumption
- correct behavior

For each issue provide:
- file/function
- expected shape/meaning
- actual transformation
- why it is wrong
- smallest safe fix
- a targeted test using a tiny deterministic attention example

Output:

## Attention shape ledger
A step-by-step table from input through output.

## Mask semantics
State exactly which positions are allowed and forbidden.

## Findings
Ranked by severity.

## Minimal corrections
No architecture redesign unless required by a proven invariant violation.

## Verification
Include shape assertions plus at least one tiny mask test where expected attention eligibility is known in advance.