Skip to content

fix(strategy): size reconstruct() row-block from compute_dtype, not out_dtype#117

Open
boskodev790 wants to merge 1 commit into
zeokin:mainfrom
boskodev790:fix/reconstruct-row-block-compute-dtype-66
Open

fix(strategy): size reconstruct() row-block from compute_dtype, not out_dtype#117
boskodev790 wants to merge 1 commit into
zeokin:mainfrom
boskodev790:fix/reconstruct-row-block-compute-dtype-66

Conversation

@boskodev790

Copy link
Copy Markdown

Lane

fix / bug — corrects repository behavior. CPU-safe validation only; no GPU scorecard.

Problem

strategy.subspace.reconstruct() streams its (n × n) output row-by-row to stay inside a VRAM budget, sizing the row-block from the output dtype's itemsize:

blk = _row_block(n, n, backend, np.dtype(out_dtype).itemsize, frac)

But every device tensor inside the loop — Q, Ctil, QT, and the (rb, n) product from backend.matmul — lives in cfg.compute_dtype, not the output dtype. For dtype="fp16" those differ: compute_dtype is bumped to fp32 for accuracy (see Config.compute_dtype). So the row-block is sized for half the real per-row footprint, and on an actual GPU at --dtype fp16 with large N / tight vram_fraction the streamed path can OOM. This is the same class of under-budget as the auto_tile operand-upcast fixes, on the reconstruct path.

Fix

Size the block from the dtype the loop actually computes in — mirroring how compress() / stream_gemm_* already pass compute_dtype:

def reconstruct(Ctil, Q, C_out, backend, out_dtype, frac=..., compute_dtype=None):
    item_dtype = compute_dtype if compute_dtype is not None else out_dtype
    blk = _row_block(n, n, backend, np.dtype(item_dtype).itemsize, frac)

and pass cfg.compute_dtype at the call site. compute_dtype=None defaults to out_dtype, so the existing 5-arg call in test_subspace.py is unchanged. Only strategy/subspace.py and tests/ change — no scorer/eval/docs/CI paths.

Tests

Added tests/test_reconstruct_row_block.py (CPU-only, no GPU — pure sizing arithmetic via a fake backend + a _row_block spy):

  • fp16 output + fp32 compute → the row-block is sized from fp32 (4 bytes), not fp16 (2) — fails on the pre-fix code, passes after.
  • compute_dtype=None falls back to out_dtype (back-compat).
  • output math (Q @ Ctil @ Qᵀ) is unchanged.
uv run python -m strategy.smoke
uv run --extra test python -m pytest tests/ strategy/tests/ eval/tests/ -q
# 81 passed, 23 skipped

Fixes #66

…ut_dtype

fix / bug lane — CPU-safe validation only, no GPU scorecard.

subspace.reconstruct() streams its (n x n) output row-by-row and sized the
row-block from out_dtype's itemsize. But the loop tensors (Q, Ctil, QT and the
(rb, n) product) all live in compute_dtype, which for fp16 inputs is bumped to
fp32 for accuracy. Sizing from out_dtype therefore under-budgets the row-block
by 2x at --dtype fp16, risking OOM on large-N / tight-vram_fraction runs.

Fix: add an optional compute_dtype param (defaults to None -> out_dtype for
back-compat) and size _row_block from it; pass cfg.compute_dtype at the call
site. Mirrors how compress()/stream_gemm already size from compute_dtype.

Adds tests/test_reconstruct_row_block.py (CPU-only, no GPU): asserts the
row-block is sized from compute_dtype (fp32) when out_dtype is fp16, that
compute_dtype=None falls back to out_dtype, and that the math is unchanged.

Validation: uv run python -m strategy.smoke; uv run --extra test python -m
pytest tests/ strategy/tests/ eval/tests/ -q  (81 passed, 23 skipped).

Fixes zeokin#66
@github-actions github-actions Bot added area:strategy Smart strategies / transforms (strategy/) area:tests Correctness gates and test suites (tests/) status:needs-review Awaiting maintainer review type:bug Something is incorrect or broken status:ready-non-gpu Fix/docs PR cleared non-GPU triage; maintainer review only labels Jul 8, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:strategy Smart strategies / transforms (strategy/) 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.

[bug] reconstruct() sizes its streaming row-block from out_dtype instead of the compute_dtype it actually runs in (2x under-budget for --dtype fp16)

1 participant