mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-01-09 15:08:02 -05:00
schedule: cache unbinds for consistent cache keys (#13662)
* schedule: cache unbinds for consistent cache keys different bound variable values (e.g. kv cache positions) now produce the same schedule cache key by unbinding BIND(DEFINE_VAR, CONST) before computing the cache key and rebinding after lookup. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com> * schedule: cache unbinds for consistent cache keys When scheduling, BIND(DEFINE_VAR, CONST) nodes are now unbound to tagged DEFINE_VARs before computing the cache key. This ensures that the same computation with different bound values (e.g., different KV cache positions in LLM) gets the same cache key and reuses the cached schedule. The fix: - pm_pre_sched_cache: replaces BIND with tagged DEFINE_VAR - pm_post_sched_cache: restores tagged DEFINE_VAR back to original BIND - pm_remove_rangeify_tags: excludes DEFINE_VAR to preserve tags through rangeify - var_vals extracted from BINDs before cache key computation 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com> * schedule: fix BIND handling and add CLAUDE.md - Handle BIND to RANGE in create_schedule (not matched by CONST pattern) - Assert all BINDs on same variable have same value - Add CLAUDE.md codebase guide 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
@@ -1,8 +1,31 @@
|
||||
import unittest
|
||||
from tinygrad import Tensor
|
||||
from tinygrad import Tensor, Variable
|
||||
from tinygrad.engine.schedule import schedule_cache
|
||||
|
||||
class TestScheduleCache(unittest.TestCase):
|
||||
def test_bound_variable_reuses_cache(self):
|
||||
schedule_cache.clear()
|
||||
v = Variable('v', 1, 100)
|
||||
x = Tensor.ones(10).contiguous().realize()
|
||||
|
||||
# first run with v=5
|
||||
t1 = (x + Tensor(v.bind(5))).sum()
|
||||
self.assertEqual(t1.item(), 60.0)
|
||||
cache_size_after_first = len(schedule_cache)
|
||||
|
||||
# second run with v=10 should reuse cache
|
||||
t2 = (x + Tensor(v.bind(10))).sum()
|
||||
self.assertEqual(t2.item(), 110.0)
|
||||
self.assertEqual(len(schedule_cache), cache_size_after_first)
|
||||
|
||||
def test_bound_variable_var_vals(self):
|
||||
v = Variable('pos', 1, 100)
|
||||
x = Tensor.ones(10).contiguous().realize()
|
||||
|
||||
t = x + Tensor(v.bind(42))
|
||||
_, var_vals = t.schedule_with_vars()
|
||||
self.assertEqual(var_vals, {'pos': 42})
|
||||
|
||||
def test_simple(self):
|
||||
a = Tensor.ones(10).contiguous()
|
||||
b = Tensor.ones(10).contiguous()
|
||||
|
||||
Reference in New Issue
Block a user