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: