Repository navigation
Conversation
…und's plan and hidden rows (Linux and macOS) zig/src/fabric/stage is the first step of ashhart#2005: each rank listens on its own address only, connects to the next rank and accepts the one before. A hello checks that both ends of every link agree on the world, on being neighbours, and on the checkpoint's and layout's fingerprint. Frames carry a 32-byte header (magic, version, kind, round, rows, payload bytes) and every send and receive has a deadline. A frame of another kind or round, or past the receiver's buffer, is refused, and any failed call leaves the link broken, since its stream may be mid-frame. A rank that fails closes both links, so its neighbours see it and the ring ends as one. It uses libc sockets only, with no CUDA or Metal module. fabric/tp's control plane imports cuda through config.zig, so this does not reuse it. TF_STAGE_RANK, TF_STAGE_WORLD, TF_STAGE_PEERS and TF_STAGE_TIMEOUT_S name a rank's settings. `zig build test-stage` runs rings of two and three ranks on loopback threads, and `test-stage-compile` builds the tests for aarch64 Linux and macOS.
… any frame A stage rank opens a port, so a ring needs a key: TF_STAGE_SECRET_FILE names a file of at least 16 bytes that only its owner may read, and the key is its SHA-256. When a link forms, both ends send a fresh nonce from /dev/urandom, then an HMAC-SHA256 over a label, their role (connector or acceptor), rank, world, the fingerprint and both nonces. Each end checks the other's proof in constant time, so the key never crosses the wire and a proof can't be replayed. A rank without the key is refused by both its neighbours (error.BadSecret), and a stranger that connects and talks first is refused at its first frame. This authenticates the ring when it forms. It doesn't encrypt, and it doesn't protect each frame after that, so it's for a trusted network.
…wildcard listen and a key file over 4 KiB are refused
…y report, rank 0's round plan, the rows layout and the keep-alive, each mismatch refused by its name
…yte-stable and each refused by its name, for a CUDA or Metal stage
…ytes on the wire read TFSP; fixtures rewritten
…d the keep-alive beats mid-round too RFC ashhart#2005 section 6.1 has a stage refuse rows whose generation is not its own, but the plan carried none. A stream is now 28 bytes: the generation u32 at 24, 0 for an open (BadOpen otherwise), and plan.sameGeneration refuses another by name (OtherGeneration). The keep-alive is due once 2 s pass without a plan, not only when idle, so a stage's long forward elsewhere in the ring never reads as a dead ring; a stage already counts any frame as heard. Ready's warm means executable for the declared shapes, not graphs captured. The plan fixtures are rewritten (140 bytes for the round) and plan_open_generation added.
|
Thanks @jschmied, byte-equal logits at 2 and 3 stages with both KV formats is a strong start for the shape, and the memory split is what the RFC was after. Our answers, in order: put the mode byte into v0, in #2060, before v0 freezes, since Metal stages need it too; boot and plan echoes around the ring are fine for v1, noted as a limit in section 6.4; the idle-only keep-alive is fine for this step; a kept prefix is enough for now, with tree paths as their own protocol step; and refusing resume and drafter taps by name until steps 7 and 10 is right. Rebase once #2012 and #2060 are in. |
… prompt chunk, so a stage picks the engine's kernels A Qwen stage runs a prompt chunk through the prompt GEMMs and a window through the one-row-exact path, and the bits differ, so a follower must know which one rank 0 runs. A keep-alive is a window; any other mode byte is BadMode. Fixtures plan_prompt, plan_bad_mode and plan_keepalive_mode added; plan_bad_reserved pokes byte 21.
|
Thanks @ashhart. The mode byte is in #2060 now, in RFC #2005 section 6.4 now has the round mode among v0's decisions, and a "Limits of v1" note: plan echoes around the ring, the idle-only keep-alive, kept prefixes without tree paths, and resume and taps refused until steps 7 and 10. I'll rebase this PR once #2012 and #2060 are in, and it then drops its own copy of the mode commit. |
…e passes it on to the next recv marked the link broken on any error, a received stop too, so the rank that got it could only hang up: in a ring of three the last stage saw Closed where rank 0 had stopped cleanly. A stop is a whole frame; only other errors break the link now.
… binding, so CUDA, Metal and XPU stages share them driver.zig takes the backend-neutral half of ashhart#2115's Qwen CUDA stage: the boot over the ring (descriptions, then ready reports, rank 0 checking each), the fingerprint, the slots that apply a plan's open, keep, close and generation checks alike on every rank, rank 0's rounds with the plan's echo, and a follower's serve loop. `Driver(B)` is generic over a binding that names its engine, cache and forward part and supplies the boundary, the shapes, the cache's position, window and rollback, and the forward of its layer range given the round's mode. driver_test.zig runs it over a fake engine whose every layer mixes its index and a per-stream state into the rows: 2- and 3-stage loopback rings give one engine's bytes, a stage one layer short and a rollback that moves the position but not the state both change the bytes, a dropped rollback is refused as OtherPosition, and a range described against the layout is refused at boot as OtherRange.
…es on the wire read TFST, and the frame test checks them
…to a stalled peer keeps its deadline macOS ignores MSG_DONTWAIT on a blocking socket: sendAll's send of a 64 MiB payload waited for the peer to take every byte, past the deadline, and the mid-payload stop test hung on an M3 Pro. connect no longer restores the blocking flag after it connects, and an accepted socket is made non-blocking too; the poll loops already take partial sends.
…534654 (TFST), rewritten
…nd state, and runs its part of the forward RFC ashhart#2005 step 2. Geometry.stage slices a model to layers [first, end): linear, bits and kinds stay indexed by global layer, deltas() and attentions() count the stage's own, so every state, KV, log and keep size follows. The loader names tensors by global layer and loads the embedding only from layer 0 and the final norm and head only where the stage holds them (MLX checkpoints; ModelOpt, EXL3 and pockets refused). A stage without the head allocates no logits scratch. The 27B's recurrent scratch path keys on the whole model, as one engine takes it. Engine.initStage builds such an engine. Its forward is cuda_stage_forward.run: the boundary's h and r come in (past rank 0), its layers run, h and r go out, and rank 0 then runs the final norm and head over the rows the ring brings back. Shape.part keeps the layers and head pieces in their own captured graphs; the TCP hop is outside both. The first layer of a later stage adds the incoming residual, as layer i of one engine does. One engine's path is unchanged: step, window and prompt logits equal the parent's bit for bit.
…nd tf-qwen35-cuda stage cuda_stage.zig drives the stage link (ashhart#2012) with the protocol (ashhart#2060): at boot each stage's description and ready report go round the ring and rank 0 checks every one and the layout; each round rank 0's plan goes round (rank 0 checks the last stage's echo byte for byte), then the rows, h and r as bf16 planes. Slots is the same code on every rank: open takes a free slot at generation 0, continue keeps the previous window's first `accepted` rows (a keep advances the generation) and checks the plan's generation and position, close frees it; resume is refused until the prompt cache spans stages. Rank 0's Lead predicts each entry's position and generation, sends the plan, then applies it through Slots like every follower. tf-qwen35-cuda <model> stage <out> runs one scenario (a two-chunk prompt, greedy steps, a second stream beside it, drafted windows kept to 3, 5, 0 and 2 rows, a keep-alive) on one engine or a ring. Qwen3.8-27B MLX 4-bit on one GB10: 2 stages (0:32, 32:64) and 3 (0:20, 20:44, 44:64) give the one engine's 48 logits rows byte for byte, bf16 and fp8 KV; device bytes 16.98 GiB whole, 9.04 + 7.64 GiB in two, 6.61 + 6.03 + 5.22 GiB in three. A follower killed mid-scenario ends rank 0 with Closed 0.4 s later.
…es built field by field are one engine The engine's checkpoint scope moved into initStage with the stage span. Test geometries that start undefined set first and whole to 0, or deltas() reads past the layer table.
…ot and rounds now shared with other backends cuda_stage.zig keeps only what is Qwen's: the version, the boundary (h and r, bf16), the shapes, the KV dtype, a slot's cache, its position, window and prefix-path rollback, and the forward of the stage's layers with the round's hop, a prompt chunk keeping its last row's logits and a window every row's. The boot, the slots, rank 0's rounds and the follower's serve loop are the driver's. Lead.step takes the plan's mode instead of the engine's keep.
305e021 to
46126d0
Compare
jschmied
changed the title
|
I split this PR in two, so that Metal and XPU stages don't have to copy the Qwen code:
The logits are still byte-equal to the whole model at 2 and 3 stages, with bf16 and fp8 KV, on the new head |
|
The speed beside the tokens, as #2005 asks: Qwen3.8-27B (MLX 4-bit) on one GB10, the single engine against 2 stages (layers 0:32, 32:64) as two processes on loopback. Arms C and D shape the ring's TCP ports with netem on The binary is v1 Setup: bf16 KV, 4 slots of 33280 rows, real text as prompts, 4096-row prompt chunks. Each arm gets a fresh process and the page cache is dropped before each start. A and B ran 3 rounds (AB, BA, AB). C and D ran 1 round each. Each cell shows the median and the min-max over the rounds.
Tokens: all 8 runs wrote the same bytes. That covers the first-token logits of 7 prompts and 650 greedy tokens from 5 streams (1 alone, 4 side by side). Shaping the link changes no token. Decode is bound by latency. On loopback the split costs 0.7 % (14.85 to 14.74 tok/s). Over a link, each round adds about two one-way delays: +22.9 ms a step at 10 ms a hop and +45.2 ms at 20 ms. With 4 streams the same delay is shared by 4 rows (52.6 to 39.6 and 30.4 tok/s). Long prompts are bound by bandwidth. v1 sends every row back from the last stage, so a prompt crosses the link twice at 20 KiB a row. At 100 Mbit the 32k prompt spends 115 s on the wire (144.0 s against 29.0 s on loopback), and half of that is the return of rows the head does not use. Sending only the last row back for a prompt chunk would halve the wire time. One more cost: on plain loopback the split adds about 11 % to TTFT (26.2 to 29.0 s at 32k), about 0.17 s for each 4096-row hop. I have not split that cost into its parts yet. |
RFC #2005 step 2: Qwen3.8-27B (
qwen3_5, CUDA) as a pipeline stage. Stacked on #2241 (the stage driver), which is stacked on #2060, which is stacked on #2012. This PR is a draft for now. It needs a rebase once its parents land, and ATC reports conflicts with several openqwen3_5PRs.What changed since v1 (
305e0219b): the stage driver moved out into #2241, so that Metal and XPU stages can reuse it.cuda_stage.zigis nowfabric/stage/driver.zig, asDriver(B)over an engine binding. That part is the boot, the slots, rank 0's rounds and the follower's serve loop.cuda_stage.zig(325 → 84 lines) is now only the Qwen CUDA binding.Commits
b1ac372f5,cb4798211) and its fixtures (3a6b096ed). They fall away when fabric stage: the protocol layer, boot description, round plan, rows layout and keep-alive, with golden bytes for every message (#2005 step 1b, stacked on #2012) #2060 rebases on fabric: the pipeline stage link, a TCP ring that carries each round's plan and hidden rows, with a shared key (Linux and macOS, #2005 step 1) #2012.fabe56ddf: a stage holds and loads only its layer range;c5771036e: the engine as a stage, andtf-qwen35-cuda <model> stage;9e509fc48: the load-scope check;46126d072: the stage as a binding of the driver.What it does
cuda_stage.zig) gives the driver:cuda_weights.zig,cuda_engine.zig,cuda_geometry.zigandcuda_shapes.zig, unchanged from v1. The split forward is incuda_stage_forward.zig. ItsHopstays the engine's own type, so the engine still builds without thestagemodule; the binding converts the driver's hop to it.tf-qwen35-cuda <model> stage.Lead.stepnow takes the plan's mode (prompt or window) instead of the engine's keep.Results (one GB10, Qwen3.8-27B MLX 4-bit, binary
b78652dc18c751d317ff3ff6287a2abeat46126d072)The scenario covers:
Each comparison covers 48 logits rows, 23,838,720 bytes.
305e0219b)Every follower exits 0. Device memory is the same as in v1:
Checks:
zig build test: RC 0, 0 errorsQuestions
acceptedmeans a kept prefix of the window. Keeping a tree path would need the path in the plan.