github-starred/cs249r_book
2
0
Fork 0
mirror of https://github.com/harvard-edge/cs249r_book.git synced 2026-07-17 08:28:07 -05:00
Code Issues 276 Packages Projects Releases 26 Wiki Activity

[PR #1875] fix(attention): mask reshape assumes square (seq_q==seq_k), breaks cross-attention #27336

New Issue
Open
opened 2026-06-18 16:55:34 -05:00 by GiteaMirror · 0 comments
GiteaMirror commented 2026-06-18 16:55:34 -05:00
Owner
Copy Link

📋 Pull Request Information

Original PR: https://github.com/harvard-edge/cs249r_book/pull/1875
Author: @Shashank-Tripathi-07
Created: 6/16/2026
Status: 🔄 Open

Base: main ← Head: fix/attention-mask-reshape-cross-attention

📝 Commits (1)

  • 2544b54 fix attention mask reshape to preserve both seq_q and seq_k dimensions

📊 Changes

1 file changed (+4 additions, -3 deletions)

View changed files

📝 tinytorch/src/12_attention/12_attention.py (+4 -3)

📄 Description

What breaks

MultiHeadAttention.forward() reshapes a 3-D attention mask (batch, seq_q, seq_k) to (batch, 1, seq_q, seq_k) for multi-head broadcasting. The old code:

batch_size_mask, seq_len_mask, _ = mask.shape   # discards seq_k!
mask_reshaped = mask.reshape(batch_size_mask, 1, seq_len_mask, seq_len_mask)

The _ discards seq_k. The reshape then uses seq_len_mask (which is seq_q) for both the query and key dimensions. For self-attention where seq_q == seq_k, this accidentally works. For cross-attention (e.g. encoder-decoder attention, or attention with a different context length), seq_q != seq_k, so the reshape either:

  • Raises ValueError if seq_q * seq_q != seq_q * seq_k (total elements don't match)
  • Silently reinterprets memory if they happen to share a common product

Fix

batch_size_mask, seq_q_mask, seq_k_mask = mask.shape
mask_reshaped = mask.reshape(batch_size_mask, 1, seq_q_mask, seq_k_mask)

Unpack both dimensions explicitly. Correct for square and rectangular masks.

Test plan

  • pytest tinytorch/tests/12_attention/ passes
  • Self-attention with square mask (B, T, T) still works
  • Cross-attention with rectangular mask (B, T_q, T_k) where T_q != T_k now works

🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.

## 📋 Pull Request Information **Original PR:** https://github.com/harvard-edge/cs249r_book/pull/1875 **Author:** [@Shashank-Tripathi-07](https://github.com/Shashank-Tripathi-07) **Created:** 6/16/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `fix/attention-mask-reshape-cross-attention` --- ### 📝 Commits (1) - [`2544b54`](https://github.com/harvard-edge/cs249r_book/commit/2544b54071be33706d1b74cd84c60c0947b83061) fix attention mask reshape to preserve both seq_q and seq_k dimensions ### 📊 Changes **1 file changed** (+4 additions, -3 deletions) <details> <summary>View changed files</summary> 📝 `tinytorch/src/12_attention/12_attention.py` (+4 -3) </details> ### 📄 Description ## What breaks `MultiHeadAttention.forward()` reshapes a 3-D attention mask `(batch, seq_q, seq_k)` to `(batch, 1, seq_q, seq_k)` for multi-head broadcasting. The old code: ```python batch_size_mask, seq_len_mask, _ = mask.shape # discards seq_k! mask_reshaped = mask.reshape(batch_size_mask, 1, seq_len_mask, seq_len_mask) ``` The `_` discards `seq_k`. The reshape then uses `seq_len_mask` (which is `seq_q`) for both the query and key dimensions. For self-attention where `seq_q == seq_k`, this accidentally works. For cross-attention (e.g. encoder-decoder attention, or attention with a different context length), `seq_q != seq_k`, so the reshape either: - Raises `ValueError` if `seq_q * seq_q != seq_q * seq_k` (total elements don't match) - Silently reinterprets memory if they happen to share a common product ## Fix ```python batch_size_mask, seq_q_mask, seq_k_mask = mask.shape mask_reshaped = mask.reshape(batch_size_mask, 1, seq_q_mask, seq_k_mask) ``` Unpack both dimensions explicitly. Correct for square and rectangular masks. ## Test plan - [ ] `pytest tinytorch/tests/12_attention/` passes - [ ] Self-attention with square mask `(B, T, T)` still works - [ ] Cross-attention with rectangular mask `(B, T_q, T_k)` where `T_q != T_k` now works --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>
GiteaMirror added the pull-request label 2026-06-18 16:55:34 -05:00
Sign in to join this conversation.
No Branch/Tag Specified
Branches Tags
Labels
Clear labels
area: book
area: collabs
area: kits
area: labs
area: socratiq
area: staffml
area: tinytorch
area: tools
area: website
bug
dependencies
format: epub
format: pdf
javascript
link-health
link-rot
priority-high
pull-request
Mirrored from GitHub Pull Request
 staffml
type: bug
type: citation
type: code
type: errata
type: improvement
type: new
type: question
vault-sli
No Label pull-request
Milestone
No items
No Milestone
Projects
Clear projects
No project
Assignees
Clear assignees
GiteaMirror ninjasurge
No Assignees
1 Participants
Notifications
Due Date
No due date set.
Dependencies

No dependencies set.

Reference: github-starred/cs249r_book#27336
Delete Branch "%!s()"

Deleting a branch is permanent. Although the deleted branch may continue to exist for a short time before it actually gets removed, it CANNOT be undone in most cases. Continue?