From 4afca0e10a2083169c1a3f6c21abd003ee8f5c9f Mon Sep 17 00:00:00 2001 From: RenzoMXD <170978465+RenzoMXD@users.noreply.github.com> Date: Tue, 14 Jul 2026 07:11:56 +0200 Subject: [PATCH] fix(matmul): budget the transient fifth tile in _gemm_tiled_sync's per-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. --- matmul/gemm.py | 22 ++++++++++----- tests/test_auto_tile_budget.py | 49 ++++++++++++++++++++++------------ 2 files changed, 47 insertions(+), 24 deletions(-) diff --git a/matmul/gemm.py b/matmul/gemm.py index 067a096..abd28b8 100644 --- a/matmul/gemm.py +++ b/matmul/gemm.py @@ -36,14 +36,22 @@ def _tile_operand_bytes(cfg: Config) -> int: def _tile_workspace_bytes_per_elem(cfg: Config) -> int: """Device bytes per T×T element in one tiled (i, j, k) accumulation step. - Four T×T tiles are live at the peak of each k-step: the accumulator, the two - operand tiles (A, B), and the freshly allocated GEMM output returned by - ``backend.matmul`` — ``torch.bmm`` cannot write its result into either - operand, so ``acc``, ``a_dev``, ``b_dev`` and ``prod`` all coexist. The - output tile is produced in the operand/compute dtype, so it costs one more - ``_tile_operand_bytes`` term. + Four T×T tiles are steady-state live at each k-step -- the accumulator, + the two operand tiles (A, B), and the freshly allocated GEMM output + returned by ``backend.matmul`` (``torch.bmm`` cannot write its result into + either operand) -- but ``_gemm_tiled_sync`` reassigns ``a_dev``, ``b_dev``, + and ``prod`` every iteration (e.g. ``a_dev = backend.to_device(...)``), and + Python evaluates the right-hand side of an assignment -- allocating the NEW + tile -- before the name rebinds and the OLD tile's reference is dropped. So + at the moment any one of those three lines runs, the previous iteration's + tile it is 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, risking OOM at a ``vram_fraction`` that leaves little headroom. + The output tile (and this fifth, transient one) are produced in the + operand/compute dtype, so each costs one ``_tile_operand_bytes`` term. """ - return cfg.acc_dtype.itemsize + 3 * _tile_operand_bytes(cfg) + return cfg.acc_dtype.itemsize + 4 * _tile_operand_bytes(cfg) def auto_tile(n: int, cfg: Config, backend: Backend) -> int: diff --git a/tests/test_auto_tile_budget.py b/tests/test_auto_tile_budget.py index 04adc24..0c11739 100644 --- a/tests/test_auto_tile_budget.py +++ b/tests/test_auto_tile_budget.py @@ -38,35 +38,50 @@ def test_tile_operand_bytes_fp16_upcast(): def test_tile_workspace_counts_fp16_upcast(): fp16_acc = Config(dtype="fp16", accumulate_fp32=True) - # acc fp32 (4) + two fp32 operand tiles (4 + 4) + the fp32 bmm output tile (4) - assert _tile_workspace_bytes_per_elem(fp16_acc) == 16 + # acc fp32 (4) + four fp32 operand-sized tiles (4*4) -- see + # test_tile_workspace_counts_the_transient_fifth_tile for why it's four. + assert _tile_workspace_bytes_per_elem(fp16_acc) == 20 fp16_raw = Config(dtype="fp16", accumulate_fp32=False) - # acc fp16 (2) + two fp16 operand tiles (2 + 2) + the fp16 bmm output tile (2) - assert _tile_workspace_bytes_per_elem(fp16_raw) == 8 + # acc fp16 (2) + four fp16 operand-sized tiles (4*2) + assert _tile_workspace_bytes_per_elem(fp16_raw) == 10 def test_tile_workspace_counts_the_bmm_output_tile(): - """Regression for #95: the per-step working set holds four T×T tiles — + """Regression for #95: the per-step working set holds a fourth T×T tile — acc + A + B + the GEMM output `prod` — so the model must budget the output - tile too. Before the fix it counted only acc + 2 operands (undersizing the - per-element cost and over-picking T).""" + tile too, not just acc + 2 operands.""" for cfg in ( Config(dtype="fp32"), Config(dtype="fp16", accumulate_fp32=True), Config(dtype="fp16", accumulate_fp32=False), ): operand = _tile_operand_bytes(cfg) - # acc + two operand tiles + one output tile (output produced in operand dtype). - assert _tile_workspace_bytes_per_elem(cfg) == cfg.acc_dtype.itemsize + 3 * operand - # It must be exactly one operand tile larger than the old 3-term model. - assert ( - _tile_workspace_bytes_per_elem(cfg) - == cfg.acc_dtype.itemsize + 2 * operand + operand - ) - - # fp32 concretely: 4 (acc) + 4 + 4 (operands) + 4 (output) = 16, not the old 12. - assert _tile_workspace_bytes_per_elem(Config(dtype="fp32")) == 16 + # Must be at least acc + three operand tiles (two operands + output) -- + # the current model charges a fourth on top; see the transient-fifth- + # tile test below for why. + assert _tile_workspace_bytes_per_elem(cfg) >= cfg.acc_dtype.itemsize + 3 * operand + + +def test_tile_workspace_counts_the_transient_fifth_tile(): + """_gemm_tiled_sync reassigns a_dev, b_dev, and prod every k-iteration + (e.g. `a_dev = backend.to_device(...)`). Python evaluates the right-hand + side -- allocating the NEW tile -- before the name rebinds and the OLD + tile's reference drops, so whichever of those three lines is executing has + the previous iteration's tile it's about to replace still live alongside + the freshly allocated one: a fifth T×T tile, momentarily, on top of the + steady-state four (acc, a_dev, b_dev, prod). Budgeting only four + under-counts true peak device usage by one operand-sized tile.""" + for cfg in ( + Config(dtype="fp32"), + Config(dtype="fp16", accumulate_fp32=True), + Config(dtype="fp16", accumulate_fp32=False), + ): + operand = _tile_operand_bytes(cfg) + assert _tile_workspace_bytes_per_elem(cfg) == cfg.acc_dtype.itemsize + 4 * operand + + # fp32 concretely: 4 (acc) + 4*4 (four operand-sized tiles) = 20, not 16. + assert _tile_workspace_bytes_per_elem(Config(dtype="fp32")) == 20 def _legacy_auto_tile(n: int, cfg: Config, free_bytes: int) -> int: