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] [CLOSED] fix(attention): mask reshape assumes square (seq_q==seq_k), breaks cross-attention #33428

New Issue
Closed
opened 2026-07-13 12:03:52 -05:00 by GiteaMirror · 0 comments
GiteaMirror commented 2026-07-13 12:03:52 -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: ❌ Closed

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

📝 Commits (10+)

  • c1531ca Merge dev into main: workflow site_only standardization + book caption work + PR #1744 (ch1 PDF layout)
  • e61c8ee Update contributors list [skip ci]
  • 6736f9d Merge dev into main: auto-tidy fix for book-format hooks on intro.qmd
  • 5f53073 Merge remote-tracking branch 'origin/main'
  • d562a28 Merge dev into main: revert failed v0.7.0 bump + contributors sync
  • 2854c1d Merge remote-tracking branch 'origin/dev'
  • f26abeb 🔥 TinyTorch tinytorch-v0.2.0: Module clarity improvements for 04_losses (loss landscape ASCII art), 05_dataloader, 06_autograd. Pyproject sphinxcontrib-mermaid version bump.
  • 2e04a1f Merge dev into main: revert v0.7.0/v0.2.0 bumps after cancelled publish
  • 2169a1e 🔥 TinyTorch tinytorch-v0.1.12: Module clarity improvements for 04_losses (loss landscape ASCII art), 05_dataloader, 06_autograd. Pyproject sphinxcontrib-mermaid version bump.
  • df88ff7 fix(ci): isolate dispatch/call runs in tinytorch+labs validate concurrency

📊 Changes

5 files changed (+229 additions, -5 deletions)

View changed files

📝 .github/workflows/labs-validate-dev.yml (+5 -1)
📝 .github/workflows/tinytorch-validate-dev.yml (+6 -1)
➕ site/newsletter/posts/2026/2026-05-19_from-the-builders-gap-to-a-builders-stack.md (+106 -0)
➕ site/newsletter/posts/2026/2026-05-22_team-newsletter-beyond-the-spiral-tinyml-pedagogy-at-the-hardware-edge.md (+108 -0)
📝 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:** ❌ Closed **Base:** `dev` ← **Head:** `fix/attention-mask-reshape-cross-attention` --- ### 📝 Commits (10+) - [`c1531ca`](https://github.com/harvard-edge/cs249r_book/commit/c1531ca89e55d2de66badd6125f4ae80bc03e9c2) Merge dev into main: workflow site_only standardization + book caption work + PR #1744 (ch1 PDF layout) - [`e61c8ee`](https://github.com/harvard-edge/cs249r_book/commit/e61c8eeaa930712c7ee9d3edfdbba3bed8c93184) Update contributors list [skip ci] - [`6736f9d`](https://github.com/harvard-edge/cs249r_book/commit/6736f9dff3cb1fa4dc53ed654257ab1da7bdd5be) Merge dev into main: auto-tidy fix for book-format hooks on intro.qmd - [`5f53073`](https://github.com/harvard-edge/cs249r_book/commit/5f53073037b730cf6443776cb399110755e3d44d) Merge remote-tracking branch 'origin/main' - [`d562a28`](https://github.com/harvard-edge/cs249r_book/commit/d562a28abcdb7dc51c98fe8e32a642903e0471d5) Merge dev into main: revert failed v0.7.0 bump + contributors sync - [`2854c1d`](https://github.com/harvard-edge/cs249r_book/commit/2854c1d3a2336ea0a242db3de6bfcf1cbb369a1f) Merge remote-tracking branch 'origin/dev' - [`f26abeb`](https://github.com/harvard-edge/cs249r_book/commit/f26abeb6dd746ffdb8b3cd170f0692bc9815d8a5) 🔥 TinyTorch tinytorch-v0.2.0: Module clarity improvements for 04_losses (loss landscape ASCII art), 05_dataloader, 06_autograd. Pyproject sphinxcontrib-mermaid version bump. - [`2e04a1f`](https://github.com/harvard-edge/cs249r_book/commit/2e04a1f3edf81bb3ee4b9e69d9f24a971f037e20) Merge dev into main: revert v0.7.0/v0.2.0 bumps after cancelled publish - [`2169a1e`](https://github.com/harvard-edge/cs249r_book/commit/2169a1efb617832ec53fc4f055fcc7acaaf635cb) 🔥 TinyTorch tinytorch-v0.1.12: Module clarity improvements for 04_losses (loss landscape ASCII art), 05_dataloader, 06_autograd. Pyproject sphinxcontrib-mermaid version bump. - [`df88ff7`](https://github.com/harvard-edge/cs249r_book/commit/df88ff7a72141d907a8248a343ae6cc5c2384a45) fix(ci): isolate dispatch/call runs in tinytorch+labs validate concurrency ### 📊 Changes **5 files changed** (+229 additions, -5 deletions) <details> <summary>View changed files</summary> 📝 `.github/workflows/labs-validate-dev.yml` (+5 -1) 📝 `.github/workflows/tinytorch-validate-dev.yml` (+6 -1) ➕ `site/newsletter/posts/2026/2026-05-19_from-the-builders-gap-to-a-builders-stack.md` (+106 -0) ➕ `site/newsletter/posts/2026/2026-05-22_team-newsletter-beyond-the-spiral-tinyml-pedagogy-at-the-hardware-edge.md` (+108 -0) 📝 `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-07-13 12:03:52 -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#33428
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?