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 #1953] [MERGED] fix(tinytorch): cached generation step excludes current token from self-attention #34677

New Issue
Closed
opened 2026-07-14 19:28:45 -05:00 by GiteaMirror · 0 comments
GiteaMirror commented 2026-07-14 19:28:45 -05:00
Owner
Copy Link

📋 Pull Request Information

Original PR: https://github.com/harvard-edge/cs249r_book/pull/1953
Author: @Shashank-Tripathi-07
Created: 7/7/2026
Status: ✅ Merged
Merged: 7/10/2026
Merged by: @profvjreddi

Base: dev ← Head: fix/memoization-cached-generation-missing-self-attention

📝 Commits (1)

  • 267a534 fix(tinytorch): cached generation step excludes current token from self-attention

📊 Changes

1 file changed (+7 additions, -1 deletions)

View changed files

📝 tinytorch/src/18_memoization/18_memoization.py (+7 -1)

📄 Description

Summary

  • _cached_generation_step() writes the new token's K/V into the cache via update(), then calls cache_obj.get() and uses its result directly as K_all/V_all for the attention computation.
  • KVCache.get() is correctly designed and tested (test_unit_kvcache, which explicitly asserts "Before advance, get() should return empty") to return only tokens already made visible by advance() -- exactly not the token just written by update(), since advance() is meant to run once per token after all layers finish.
  • So the new token's own K/V was never included in its own attention computation. Worse: on the very first generation step (cache empty, nothing advanced yet), cache.get() returns a zero-length sequence, crashing the attention matmul with ValueError: zero-size array to reduction operation maximum which has no identity -- cached generation was broken from the first token onward, not just subtly wrong.

Fix

Concatenate the current step's own K_heads/V_heads (already computed earlier in the same function, no extra work needed) onto cache.get()'s history, instead of using the cache's history alone. KVCache itself is untouched -- it already works exactly as designed and tested.

Test plan

  • Built a stub attention module (identity projections) against the real KVCache class.
  • Reproduced the crash with the old logic on the first token: cache.get() returns shape (1,1,0,4), and the subsequent softmax/matmul crashes with ValueError: zero-size array to reduction operation maximum which has no identity.
  • Confirmed the fixed function produces the correct self-attention output ([1,0,0,0], the token attending fully to itself, the only token present) instead of crashing.
  • Confirmed history accumulates correctly across subsequent tokens after advance().
  • python -m py_compile passes on the edited file.

🔄 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/1953 **Author:** [@Shashank-Tripathi-07](https://github.com/Shashank-Tripathi-07) **Created:** 7/7/2026 **Status:** ✅ Merged **Merged:** 7/10/2026 **Merged by:** [@profvjreddi](https://github.com/profvjreddi) **Base:** `dev` ← **Head:** `fix/memoization-cached-generation-missing-self-attention` --- ### 📝 Commits (1) - [`267a534`](https://github.com/harvard-edge/cs249r_book/commit/267a53476f8bdb4141b51b97d81ded4d4e7da90a) fix(tinytorch): cached generation step excludes current token from self-attention ### 📊 Changes **1 file changed** (+7 additions, -1 deletions) <details> <summary>View changed files</summary> 📝 `tinytorch/src/18_memoization/18_memoization.py` (+7 -1) </details> ### 📄 Description ## Summary - `_cached_generation_step()` writes the new token's K/V into the cache via `update()`, then calls `cache_obj.get()` and uses its result directly as `K_all`/`V_all` for the attention computation. - `KVCache.get()` is correctly designed and tested (`test_unit_kvcache`, which explicitly asserts "Before advance, get() should return empty") to return only tokens already made visible by `advance()` -- exactly **not** the token just written by `update()`, since `advance()` is meant to run once per token after all layers finish. - So the new token's own K/V was never included in its own attention computation. Worse: on the very first generation step (cache empty, nothing advanced yet), `cache.get()` returns a zero-length sequence, crashing the attention matmul with `ValueError: zero-size array to reduction operation maximum which has no identity` -- cached generation was broken from the first token onward, not just subtly wrong. ## Fix Concatenate the current step's own `K_heads`/`V_heads` (already computed earlier in the same function, no extra work needed) onto `cache.get()`'s history, instead of using the cache's history alone. `KVCache` itself is untouched -- it already works exactly as designed and tested. ## Test plan - [x] Built a stub attention module (identity projections) against the real `KVCache` class. - [x] Reproduced the crash with the old logic on the first token: `cache.get()` returns shape `(1,1,0,4)`, and the subsequent softmax/matmul crashes with `ValueError: zero-size array to reduction operation maximum which has no identity`. - [x] Confirmed the fixed function produces the correct self-attention output (`[1,0,0,0]`, the token attending fully to itself, the only token present) instead of crashing. - [x] Confirmed history accumulates correctly across subsequent tokens after `advance()`. - [x] `python -m py_compile` passes on the edited file. --- <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-14 19:28:45 -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#34677
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?