Skip to content

core: Spark 2.5 and Dory on CUDA load with direct reads (one checkpoint.direct call each), 17-50 % faster cold - #2449

Open
jschmied wants to merge 2 commits into
ashhart:mainfrom
jschmied:direct-families
Open

jschmied wants to merge 2 commits into
ashhart:mainfrom
jschmied:direct-families

Conversation

@jschmied

@jschmied jschmied commented Nov 7, 2026 •

Copy link
Copy Markdown
Contributor

What this changes

Spark 2.5 and Dory on CUDA now read each layer past the page cache, with the next layer read while the current one packs: ledger row 22 on two more families. Each family has one call on core/checkpoint_direct.zig (#2282's Layers), which gains two helpers:

  • Layers.layerAt(fmt, i, next) enters layer i and starts layer next's read (the two prefix prints every caller wrote).
  • Layers.moved(v, fmt, i, n) does the same, then moves the tensor views inside v (fields, unions, optionals, slices) onto the layer's copies. It is for loaders that map every tensor before the layer loop. A view outside the layer keeps the mapping.

Spark 2.5 reads through Checkpoint.get, so it uses layerAt. Dory's names.map takes every view first, so it uses moved. Readers come from TF_LOAD_PREFETCH (default 16), as on Gemma 4 and Qwen 3.5. While direct reads are on, Spark's opt-in read-ahead gets 0 readers. TF_LOAD_DIRECT=0 or TF_LOAD_PREFETCH=0 keeps main's mapped reads.

Not here:

Checks on this head

  • Head commit: 151cb02
  • Lean checker on base and head: 116 / 116, 0 new.
  • Tests run: zig build test test-golden test-cuda-host: 590/590 steps, 7702 of 7765 tests passed, 63 skipped, 0 failed. test-kernel-contract test-server-security test-pull-security fuzz-cortheon, check-native-cuda-host compile-cuda-tests -Dtarget=x86_64-linux-gnu, zig fmt --check and git diff --check: all pass. The new core test covers layerAt and moved on a two-layer checkpoint (copies equal the mapping, out-of-layer views stay mapped, get returns the same copies), and the off test covers both as no-ops.
  • tf-cuda-test spark25-pages and dory-pages pass with direct reads and with TF_LOAD_DIRECT=0, and print the same RESULT lines (19 and 13).
  • CI: not run yet (no push).
  • Not run, and why: Kolibri-1, MiniMax M3 and GLM have no checkpoint on this box; discrete cards (we have only the GB10).

Receipt

  • Main against branch: one GB10, page cache dropped before every start, a warm-up start dropped, then A-B-B-A; time to /v1/models:

    model main (s) branch (s) replies concurrent 1,4 greedy
    Spark X2.6 4B 8-bit 5.03 / 5.04 2.53 / 2.53 (−50 %) identical identical
    Dory NVFP4 4.53 / 4.53 3.78 / 3.78 (−17 %) identical identical

    Spark saves more than Dory because main reads its 3.74 GiB of layer weights through faults as stage.upload copies from the mapping, at about 1.5 GiB/s, which is about 2.5 s of main's 5.0 s. The branch reads them on reader threads while the previous layer packs, so the whole read leaves the load's critical path. Dory has 1.73 GiB in layers (its embedding and head, 1.25 GiB, stay mapped), so it saves 0.75 s. Memory at ready is level (MemAvailable within 0.26 GiB, server RSS 6.14 and 7.18 GiB in both arms), and page-cache growth is 0.09–0.10 GiB in both arms.

  • Boxes it reaches: CUDA on Linux, Spark 2.5 and Dory loads.

  • New switch? None; TF_LOAD_DIRECT and TF_LOAD_PREFETCH as in core: Gemma 4 and Qwen 3.5 on CUDA load with direct reads, 28-31 % faster cold; Flash Next shares the read-ahead #2282.

  • Environment: TensorFold 964f98b against 151cb02; Zig 0.17.0; CUDA 13 nvcc, sm_121; GB10, Linux 7.0 aarch64, 128 GB unified; checkpoints: Spark X2.6 4B MLX 8-bit affine (converted locally from XHToken/Spark-X2.6-4B at 66e3ae7b), Dory NVFP4 (local, revision in SOURCE.json).

  • Paths and platforms: the load path of both families on CUDA, direct and TF_LOAD_DIRECT=0; run on sm_121 only.

  • Exactness: tools/bench_concurrent.py --tokens 128 --reps 1 --temperatures 0 --levels 1,4 --alone --serial: output hashes identical across all four arms per model.

  • Decode speed: n/a, loading only. The weights uploaded are the same bytes.

  • Prompt speed: n/a, loading only.

  • Licence: n/a.

Checklist

  • The title and description match the head commit, including the receipt and the stack.
  • New tests fail before the change, pass after it, and skip cleanly without their dependency. (The new test calls the new helpers, so it can't build on main; it skips off Linux or with TF_LOAD_DIRECT=0.)
  • Comments and docstrings are one line. No measurements or history in the source.
  • No personal data, machine names, internal hosts or local paths in the diff or its commits. No AI attribution lines.
  • Every token still goes through the lane rounds. No serial path, and nothing that needs drafts off.
  • Drafted output equals "draft": false, a resumed prompt equals a fresh one, and concurrent equals solo.
  • No precision traded for speed.
  • Decode and prompt processing are no slower than main (unchanged: loading only).

…hat maps every tensor first can read direct

layerAt enters a layer's prefix by index and starts the next one's read. moved does the same and repoints the
tensor views inside a value (fields, unions, optionals, slices) onto the layer's copies; a view outside the
layer keeps the mapping.
…hile it packs

One checkpoint.direct call per layer: layerAt in Spark 2.5's get-based loader, moved over Dory's tensor map.
While direct reads are on, Spark's opt-in read-ahead gets no readers. TF_LOAD_DIRECT=0 keeps the mapped reads.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant