Skip to content

fix(matmul): budget the transient fifth tile in _gemm_tiled_sync's per-step working set#227

Merged
zeokin merged 1 commit into
zeokin:mainfrom
RenzoMXD:feat/nystrom-iter
Jul 14, 2026
Merged

fix(matmul): budget the transient fifth tile in _gemm_tiled_sync's per-step working set#227
zeokin merged 1 commit into
zeokin:mainfrom
RenzoMXD:feat/nystrom-iter

Conversation

@RenzoMXD

@RenzoMXD RenzoMXD commented Jul 13, 2026

Copy link
Copy Markdown
Contributor

PR kind

  • fix
  • feat

(Replaces this PR's original content -- a nystrom-iter feat idea that a real GPU scorecard showed doesn't clear the dominance gate at the regime tested. Repurposing the branch/PR for unrelated, verified content instead of abandoning it.)

Summary

_tile_workspace_bytes_per_elem budgets four T×T tiles per k-step in matmul/gemm.py's tiled path: acc, a_dev, b_dev, prod. But _gemm_tiled_sync reassigns a_dev, b_dev, and prod every iteration (a_dev = backend.to_device(...), etc.), and Python evaluates the right-hand side of an assignment -- allocating the new tile -- before the name rebinds and the old tile's reference drops. So at the moment any one of those three lines runs, the previous iteration's tile it's about to replace is still live alongside the freshly allocated one: a fifth T×T tile, momentarily, on top of the steady-state four.

Budgeting only four under-counts true peak device usage by one operand-sized tile (roughly 20-25% depending on T), risking OOM at a vram_fraction that leaves little headroom -- exactly the failure mode the tiled path exists to avoid.

I verified the mechanism in isolation with a standalone CPython reference-counting trace before touching this file: steady state was 4 tile-sized objects, measured peak was consistently 5, across fp16/fp32/fp64 and accumulate_fp32 on/off.

Fix

_tile_workspace_bytes_per_elem: cfg.acc_dtype.itemsize + 3 * _tile_operand_bytes(cfg)... + 4 * _tile_operand_bytes(cfg). Budget-only change -- the blocked GEMM math (_tiles, the k-loop accumulation) is untouched.

Validation

  • uv run python -m strategy.smoke
  • uv run --extra test python -m pytest tests/ strategy/tests/ eval/tests/ -q — 288 passed, 24 skipped
  • Updated tests/test_auto_tile_budget.py's two tests that hardcoded the old 3-operand-term formula (test_tile_workspace_counts_fp16_upcast, test_tile_workspace_counts_the_bmm_output_tile), and added test_tile_workspace_counts_the_transient_fifth_tile explaining and asserting the new 4-operand-term model directly.

No protected paths touched (matmul/, tests/ are Open per CONTRIBUTING.md).

Checklist

  • CPU-safe validation passed
  • My commits do not include Co-authored-by footers for coding agents
  • Correctness gates pass

@github-actions github-actions Bot added area:strategy Smart strategies / transforms (strategy/) area:tests Correctness gates and test suites (tests/) area:docs README, CONTRIBUTING, BENCHMARKS and other docs status:needs-review Awaiting maintainer review status:needs-scorecard Feature / strategy PR is missing the reproducible eval scorecard type:feature Feature / improvement PR that must carry a GPU scorecard type:strategy New smart strategy / transform proposal labels Jul 13, 2026
…r-step working set

_tile_workspace_bytes_per_elem budgeted four T*T tiles per k-step: acc,
a_dev, b_dev, prod. But _gemm_tiled_sync reassigns a_dev, b_dev, and prod
every iteration (a_dev = backend.to_device(...), etc.), and Python
evaluates the right-hand side of an assignment -- allocating the NEW
tile -- before the name rebinds and the OLD tile's reference drops. So at
the moment any one of those three lines runs, the previous iteration's
tile it's about to replace is still live alongside the freshly allocated
one: a fifth T*T tile, momentarily, on top of the steady-state four.
Budgeting only four under-counts true peak device usage by one
operand-sized tile (roughly 20-25% depending on T), risking OOM at a
vram_fraction that leaves little headroom -- exactly the tiled path's
job to avoid.

Verified the mechanism in isolation with a standalone CPython refcount
trace before touching this file: steady state was 4 tile-sized objects,
measured peak was 5, consistently, across dtype combinations.

Budget-only change -- the blocked GEMM math (_tiles, the k-loop
accumulation) is untouched.
@RenzoMXD RenzoMXD force-pushed the feat/nystrom-iter branch from c48c68a to 4afca0e Compare July 14, 2026 05:12
@github-actions github-actions Bot added area:matmul Exact engine and tile schedules (matmul/) and removed area:strategy Smart strategies / transforms (strategy/) area:docs README, CONTRIBUTING, BENCHMARKS and other docs labels Jul 14, 2026
@RenzoMXD RenzoMXD changed the title feat(strategy): nystrom-iter -- landmark-seeded power iteration for decaying-spectrum data fix(matmul): budget the transient fifth tile in _gemm_tiled_sync's per-step working set Jul 14, 2026
@github-actions github-actions Bot added type:bug Something is incorrect or broken and removed type:feature Feature / improvement PR that must carry a GPU scorecard type:strategy New smart strategy / transform proposal status:needs-scorecard Feature / strategy PR is missing the reproducible eval scorecard labels Jul 14, 2026
@RenzoMXD RenzoMXD marked this pull request as ready for review July 14, 2026 05:13
@github-actions github-actions Bot added the status:ready-non-gpu Fix/docs PR cleared non-GPU triage; maintainer review only label Jul 14, 2026

@zeokin zeokin left a comment

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Budgets the transient 5th GEMM tile in _gemm_tiled_sync; fixes a real peak-VRAM undercount.

@zeokin zeokin merged commit 6346292 into zeokin:main Jul 14, 2026
9 of 12 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:matmul Exact engine and tile schedules (matmul/) area:tests Correctness gates and test suites (tests/) status:needs-review Awaiting maintainer review status:ready-non-gpu Fix/docs PR cleared non-GPU triage; maintainer review only type:bug Something is incorrect or broken

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants