Skip to content

st: retain tensor shapes and shard sizes - #802

Open
DrewZt wants to merge 2 commits into
JustVugg:devfrom
whale-agent-lab:codex/st-shape-rank-sizes
Open

st: retain tensor shapes and shard sizes#802
DrewZt wants to merge 2 commits into
JustVugg:devfrom
whale-agent-lab:codex/st-shape-rank-sizes

Conversation

@DrewZt

@DrewZt DrewZt commented Aug 3, 2026

Copy link
Copy Markdown

Summary

  • retain each safetensors tensor's rank and shape (up to ST_MAX_RANK=8)
  • retain the indexed size of each primary shard alongside fds[] and paths[]
  • reuse the recorded shard size during indexing instead of discarding the fstat result
  • reject non-numeric, fractional, negative, overflowing, or over-rank shapes before storing them

Why

st_tensor currently keeps only numel, so engines cannot validate the original tensor geometry after indexing. The shard size used for the initial bounds check is also discarded, which prevents later range diagnostics from using the same indexed fact.

This is the shared st.h prerequisite discussed on #165. It deliberately does not add the V4 engine, a per-tensor shard field, or model-specific aliases. sizes[512] costs 4 KiB once per model index.

Tests

  • test_st_shape validates non-numeric, fractional, negative, dimension-bound, product-overflow, and rank>8 refusals in subprocesses on Windows and POSIX
  • test_st_shape also accepts rank 8, scalar [] (rank=0, numel=1), and zero-element [0,4096] shapes with exact stored metadata
  • test_st_pread, test_st, test_st_mirror, and test_ue8m0 pass on Windows/UCRT64
  • portable build and all C gates pass under the stages used by make -C c check
  • Python: 288 tests passed, 32 skipped (PYTHONUTF8=1 on Windows)

Related: #165

@DrewZt
DrewZt marked this pull request as ready for review August 3, 2026 05:29
@DrewZt

DrewZt commented Aug 3, 2026

Copy link
Copy Markdown
Author

@JustVugg, could you please approve the fork workflow runs for the current 764c450 head?

This head adds the Windows/POSIX subprocess regression coverage for hostile shapes described in the PR body. I will use the actual job results, not the body, as CI evidence.

@JustVugg

JustVugg commented Aug 3, 2026

Copy link
Copy Markdown
Owner

CI approved and running — it had never started. GitHub was holding the workflow in action_required, and that state is not visible anywhere on the pull request page, so the PR looked simply unreviewed. It was not you. Four runs were stuck the same way and I have released them all; one had been waiting eleven days.

This is exactly the PR I asked for on #165, and it is scoped the way I hoped: shape[], rank, and sizes[], and nothing else. No V4 engine, no per-tensor shard, no model-specific aliases. You kept out everything that is yours and left in only what is shared, which is what makes a change to a header four engines read reviewable at all.

Two things I want to say properly while the checks run.

I was wrong about sizes[512] and you were right to keep it. I claimed it would cost ~290 MB, having assumed it sat in st_tensor — one per tensor across 72,317 of them. It is in shards, next to fds[]/paths[]: 4 KB once per model index. I withdrew that on #165 and it stands withdrawn here.

Retaining fsz matters more than it looks. st_init_multi already fstats every shard for its bounds check and then discarded the result, so nothing downstream could repeat that check — not at a pread, not in a diagnostic. That is not hypothetical: a truncated download reached us here (720 MB of a 3.57 GB shard, the tool reported success) and was caught once, at index time, by luck of ordering. Keeping the size makes that check repeatable.

The shape validation is the half I did not ask for and would have missed. Refusing non-numeric, fractional, negative, over-rank and product-overflowing dimensions before storing them closes a path where a hostile or corrupt mirror could put an absurd geometry into the index and have it believed by whatever reads it later. Accepting rank 8, scalar [] and zero-element [0,4096] as valid — with the stored metadata asserted exactly — is the part that shows the boundary was thought about rather than clamped.

I will watch Windows in particular: your shape tests drive exit(1) paths through subprocesses, and that platform is where the last three PRs in this repo actually broke. If it is green there, this merges.

One note for after: your branch already carries its own st.h work for V4 with a different dtype numbering. Once this lands, rebase onto it — the ordinals in dev (F8_E4M3=4, F8_E8M0=5, I64=6) are the ones to match, and since they are internal codes no container carries on disk, it is a find-and-replace on your side. Sorry for the churn; that one is on us for not asking what you already had before writing ours.

@JustVugg JustVugg added enhancement New feature or request quality Qualità del modello / quantizzazione labels Aug 3, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request quality Qualità del modello / quantizzazione

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants