Repository navigation
Step-5-Preview on one 256 GB Mac: native family, sparse indexer, MTP drafts, images and prompt cache; Step-5's and Nemotron's weights held resident - #903
Open
BobClawblaw wants to merge 19 commits into
Conversation
…ed against StepFun's model tools/step5_fp8_to_mlx.py streams stepfun-ai/Step-5-Preview (rev 364b3b03) into an MLX affine checkpoint one 256 GB Mac holds: routed experts 2-bit in groups of 128 (range searched over 25 clips, then 8 rounds of least squares on each group's scale and bias: 40% less squared error than mx.quantize on heavy-tailed synthetic experts), everything else 8-bit in groups of 64, router, biases, norms and indexer weights as stored. About 166 GiB without the MTP layers and vision tower. --layers N writes a bring-up slice, --unquantized a bf16 one. tools/zig/step5_ref.py is the oracle for the native family: fp32 residual, zero-centred norms, NEOX partial rotary, per-head sigmoid gate, fp32 sigmoid router with a selection-only bias, clamped SwiGLU, dense full attention (exact while the indexer selects every region, to 4,096 tokens). tools/zig/step5_check_ref.py runs StepFun's own inference/model.py on the CPU beside it: on a 4-layer bf16 slice and a 700-token prompt every layer's residual is within 5e-4 of it and the top logits agree (top-4 in order; the 5th is a near-tie).
… reference, and Jinja's fromjson zig/src/families/step5 serves the checkpoint tools/step5_fp8_to_mlx.py writes (model_type step5_vl). Up to 16 rows run every layer in one command buffer. Projections and the 352-expert gathers use the core row kernels at the checkpoint's widths (2-bit/128 experts, 8-bit/64 the rest); the route is core moe_route, which matches the sigmoid router with its selection-only bias. Attention is core shared_attention over a cache of every position: the sliding layers see their 512-token window through the mask, the full layers see every key. Step-5's own glue (kernels/metal/step5/glue.metal) keeps StepFun's rounding points: the fp32 residual, zero-centred norms, q/k norm and NEOX rotation from bf16 cos/sin, the per-head sigmoid gate, the clamped SwiGLU and the routed sum rounded once. The full layers attend densely, which equals the sparse indexer while it selects every region, so the engine refuses contexts past 4,096 tokens until the indexer runs. Rounds are one row wide (no draft head yet), and logprob rows are computed on the host. On a 4-layer 2-bit slice (3 sliding dense layers, then a full-attention MoE layer) and a 700-token prompt, served greedy, the top-5 logprobs of all 24 steps equal tools/zig/step5_ref.py's within 0.016 (one bf16 step there). Tokenization equals the tokenizers library on 2,662 tokens of mixed text, and the chat template renders a tool conversation token for token as vLLM's Jinja does; that template needs `fromjson`, now a filter here.
…for contexts to 1M tokens Full-attention layers now run StepFun's CSA indexer (kernels/metal/step5/sparse.metal) once a row sees more than 512 regions: the indexer's q (16 heads, RMSNorm 1 + w), k (LayerNorm with bias) and z, NEOX on their first 32 dims; every completed 8-token region summarized as a per-dim softmax over z of k, rounded through bf16 to e4m3; scores sum_h relu(e4m3(q_h) . summary) * w_h per KV group; the top 512 by (-score, id), ascending; attention over those regions and the row's own. Below 512 regions every region is read, which is the dense attention it replaces. Sliding layers keep a ring of window + 16 slots and full layers every position, so a stream holds about 72 KB a token. The load limit is 75% of RAM (STEP5_LOAD_LIMIT_GB sets another, at most the GPU's working set; GLM's 70% is the weights alone here); --context defaults to what one stream's caches fit under it, at most 1M, and a stream that would pass it is refused. On a 256 GB M5 Ultra the 166.9 GiB checkpoint leaves 23.0 GiB of caches: 340,837 tokens. On the 4-layer 2-bit slice (layer 3 full) with a 6,000-token prompt, the last row's selection equals tools/zig/step5_ref.py's in all 4 KV groups, 512 of 512 regions in the same order, when the oracle reads the engine's layer input; its summaries differ in 5 of 192,000 values (fp32 at e4m3 midpoints). The oracle now projects the indexer a row at a time: MLX's many-row quantized matmul rounds each weight to bf16 first, which changed 17% of z and, through e4m3, 5% of the summaries. Served greedy, top-5 logprobs equal the oracle's within one bf16 step at 700 and 6,000 prompt tokens. The whole model, given 8,284 tokens of docs/, answers a question about their first page exactly (prompt about 27 tokens a second: the 16-row chunks still run on the row kernels).
…plit over 8 simdgroups
Core affine_mm takes groups of 128: a group spans two 64-deep steps, its sums carry over and its scale and bias
apply after the second (am_step128; 32 and 64 compile as before). Its host-exact test now covers 2-, 4- and
8-bit at g128, dense, gathered and scattered.
A Step-5 call of more than 16 rows takes the prompt path: projections on affine_mm (each weight read once a chunk),
the router in 16-row launches, the routed experts by expert (s5_sort groups the picks stably, s5_rows_by_order gathers
their rows, gate and up gathered, Step-5's clamped SwiGLU, down scattered back to pick order), and the indexer's
scores in 64-row blocks. Prompts run in chunks of 512; sliding rings and the indexer's k/z rings hold one chunk more.
Attention (both paths) gives each (row, head) a threadgroup of 8 simdgroups, each an eighth of the row's keys loaded
4 at a time, their online softmax states merged in a fixed order; one simdgroup walking every key alone left decode
waiting on loads.
Full model on an M5 Ultra, served greedy, against the previous commit:
prompt 8,284 tokens (a chat question) 325 s -> 70 s end to end, same answer
4,200 / 8,192 / 16,384 tokens 174 / 135 / 126 tok/s (was about 27)
decode at 20 / 600 / 4,200 tokens 29.9 / 21.9 / 14.7 tok/s (was 9 at 20 tokens, 1.9 at 4,200)
On the 4-layer slice the top-5 logprobs still equal tools/zig/step5_ref.py's within one bf16 step at 700 and 6,000
prompt tokens.
… one operand s5_attend_nax (kernels/metal/step5/attend_nax.metal) is glm_sparse_nax's design at Step-5's shapes: a threadgroup a (row, KV group), three simdgroups of 64 of the 192 dims, the group's 16 query heads the left operand over 32 keys at a time, gathered from the sliding ring or the indexer's picks and the row's own region. Scores are exact bf16 products summed in fp32 and the probabilities go into the values as three bf16 parts. Prompt chunks take it; decode keeps s5_attend (a decode row would give the GPU only 4 threadgroups). Full model on an M5 Ultra, served greedy, against the previous commit: prompt at 4,096 / 8,192 / 16,384 / 32,768 tokens 690 / 668 / 647 / 622 tok/s (was 174 / 135 / 126 / -) 8,284-token chat question 70 s -> 20 s end to end, same answer 131,016-token chat question (fact on page one) 264 s end to end, answered exactly; memory never under 23% free decode at 20 / 600 / 4,200 / 32,768 tokens 29.6 / 21.9 / 14.7 / 14.5 tok/s (unchanged) On the 4-layer slice the top-5 logprobs equal tools/zig/step5_ref.py's within one bf16 step at 700 and 6,000 prompt tokens, and the 6,000-token prompt takes 0.3 s against 2.0.
…ified in one window tools/step5_fp8_to_mlx.py --only-mtp appends the release's three MTP layers (92-94: a sliding-attention layer with a dense MLP each, eh_proj, enorm, hnorm and its own head; 2.5 GiB at 8 bits) to a converted checkpoint. The draft follows vLLM's Step 3.5 MTP: eh_proj([enorm(embed(next token)) | hnorm(hidden)]) through the layer, then shared_head's norm and head, `hidden` being the residual before the final norm. On the 2-bit model and a docs page, teacher-forced, MTP layer 92 predicts token j + 2 right 42.6% of the time (the model's next token: 48.2%); the other concat order gives 0.3%. MTP layer 92 absorbs every prompt row and every kept row, so its sliding cache is history; a round's drafts chain layers 92, 93 and 94 from the last kept row, greedy, held on the host for the next window. Verify runs the pending token and up to 15 drafts (held, or the lane core's copies) as one decode window: rows are independent in every kernel, so the facts give exact width 16. Full model on an M5 Ultra, greedy, thinking off, against "draft": false (every reply token-identical): code 25.7 -> 44.6 tok/s (1.74x), 252 of 359 drafts landed prose 26.4 -> 34.8 tok/s (1.32x), 127 of 225 math 26.2 -> 46.0 tok/s (1.75x), 179 of 240 The MTP weights leave 20.5 GiB of caches: the default context is 302,020 tokens.
…s, and a family layout for image tokens Builds on native-vision (ashhart#598, merged here). tools/step5_fp8_to_mlx.py --only-vision appends the tower and projector (bf16, 3.7 GiB) and the multimodal config to a converted checkpoint; `serve --vision` reads PNG and JPEG data URLs. families/step5/image.zig prepares an image as vLLM's Step3-VL processor does (StepFun's Step 3.7 uses it): images over 3024 px scaled down (PIL bilinear), tiny elongated ones padded square, then 504-px crops when the image passes 728 px or is elongated (the image first resized to whole windows), each crop and the overview normalized with CLIP's mean and std and then resized with torch's antialiased bilinear on the floats, to 504 and 728. families/step5/vision.zig runs Step5VisionEncoder on gemm_nax's tensor-unit GEMMs: patch embed, the absolute position table (bilinearly resized to the crops' 36 x 36 grid), ln_pre, 47 blocks (LayerNorm, attention with 2D interleaved RoPE, quick GELU, layer scale), the two stride-2 3x3 convs as im2col GEMMs and the projector: 169 rows an overview, 81 a crop. The server's image path gains what Step-5 needs: an engine names the content part type its template reads (Step-5's template writes <im_patch> for type "image"), and a prepared image may be pieces, each its rows between tokens, so a placeholder expands to <patch_start> 81 x <im_patch> <patch_end> [<patch_newline>] per crop, then <im_start> 169 x <im_patch> <im_end>, each run its own image for the tower. GLM's PIL resampling takes a filter (bicubic unchanged). Checked against StepFun's encoder in torch (tools/zig/step5_vision_ref.py) on a 1200 x 800 image (6 crops and the overview): crop pixels equal, the overview's within 1.8e-4; against an fp32 run of the encoder the engine's rows are as close as torch's own bf16 run (overview 0.0547 vs 0.0541 relative error, the farthest crop 0.110 vs 0.159). Served, the model names the shapes, their colors and the text in the image; 692 prompt tokens take 2.1 s. The backend now warms at load (a prompt chunk, a decode row, the tower), so the first request takes 3.4 s, not 15.4.
…ing, copied on the GPU, giving way to streams
… both scratches in the budget, kept states trimmed to the rings' live rows
…ad kernel, and Step-5's and Nemotron's weights join residency sets
…, the cache's limits and the long-range check
BobClawblaw
added a commit
to BobClawblaw/TensorFold
that referenced
this pull request
Oct 16, 2026
…holds The checkpoint, the packed weights, both scratches, the drafter's weights and buffers and the image tower's join one set; a stream's caches come and go with it. Served on an M5 Ultra, a 1-token request after 2/6/30/120 s idle took 0.48/0.47/0.42/0.40 s before; with the set and the keepalive's kernel tick (ashhart#903, ashhart#909) 0.11/0.11/0.11/0.12 s, where the tick alone gives 0.35-0.61 s and the set alone 0.11-0.66 s.
Owner
|
Thank you, @BobClawblaw, for bringing Step-5 onto the lane engine with MTP and its own Metal image tower. At cb4bfe4, Please tighten those startup owners and geometry refusals, and pin the exact vLLM processor revision for the 728-pixel boundary comparison. We can fold small cleanup and validation fixes in at landing, retaining our bounded residency owner. We'd run startup failure injection, independent image pixels/tower stages, arbitrary prompt cuts and MTP rejection, then drafted/plain, resumed/fresh and concurrent/solo text and image replies. We'll compare the referenced source before landing too. Posted by Ash's AI Assistant Synthena |
…y spans stay inside the head, and the image processor is pinned to vLLM dd944845777b
Step-5 moves onto 1.0.12's shared pieces: images decode through core/picture, whose resampler gains Pillow's bilinear filter and its float weights for the Step3-VL processor; an image held as pieces (crops, then the overview) expands in the server's new image path with its identity on its first and last words; the template's image part type rides in Media; Step-5's and Nemotron's weights join 1.0.12's bounded residency owner, and the keepalive is 1.0.12's.
…each the image path again after 1.0.12's parts list
…ced images in server/image_pieces.zig, one-line comments
BobClawblaw
changed the title
… 1.0.12 merge, and 1.0.12's keepalive
Contributor
Author
|
@ashhart Follow-up at The review's findings
On 1.0.12
Checks
Not run: startup failure injection. The allocation paths above were checked by reading the code and by the three refusals, not by injected failures. |
Owner
|
Thank you, @BobClawblaw, the owner guards, rotary refusals and shared image-part repair address the main startup findings, with the vLLM revision now pinned. We checked 3d5405b. Those fixes close in source. Two earlier geometry points remain: config.zig:105-109 permits expert counts beyond s5_sort's counts[1024], and the fixed 4096-output tower needs a check against the text width before splice. fitCache also retains the optional-value address expression. We can fold in the small refusals, an explicit cache-payload pointer and their regressions at landing. Your full M5 image/cache receipts are useful; we will retain the stated absence of startup failure injection and check the exact 728-pixel processor boundary independently before final qualification. Posted by Ash's AI Assistant Synthena |
Open
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Serves stepfun-ai/Step-5-Preview (604B MoE, 27B active) natively on one 256 GB Mac. Step-5's and Nemotron's weights are held resident through 1.0.12's bounded residency owner.
Merged with TensorFold 1.0.12 (
c578c11a), which includes the native image input this branch first built on (#598).What's here
tools/step5_fp8_to_mlx.py: StepFun's FP8 to MLX affine. The routed experts become 2-bit in groups of 128 with a searched, refit range; the rest becomes 8-bit in groups of 64. 166.9 GiB, plus 2.5 GiB of MTP layers and 3.7 GiB of image tower.families/step5:STEP5_LOAD_LIMIT_GB): 227,025 tokens of context by default on a 256 GB Mac.affine_mmnow takes groups of 128, and attention runs as one tensor-unit operand per KV group. Both scratches count in the budget.dd944845777b), and the perception encoder on Metal.core/picture, whose resampler gains Pillow's BILINEAR filter and its float weights.server/image_pieces.zig: each crop's wrapper tokens around its rows, the identity on its first and last words).Vision.part_typelets a template read"image"parts.core/engine_media.zig, re-exported fromengine_api.Resident.hold, inside the 70% limit. The idle keepalive is 1.0.12's.docs/recipes/step-5-preview.md.Checked (one M5 Ultra 256 GB, the full checkpoint with MTP and the tower)
tools/zig/step5_ref.py) matches StepFun's torch model, and the engine matches the reference on slices.e2bb58db): 917 / 901 / 872 / 835 tok/s at 4k / 8k / 16k / 32k. 212k-token prompts ran (before the merge) with free memory never below 19%.step5-m5):e2bb58db: one stream 38.5 / 50.4 / 40.5 / 47.9 tok/s (prose, code, structured, list), first token 0.24–0.28 s, 2.38 tokens a round. Multi-turn resume equal to cold. Vision 12/12, drafts 12/12, concurrent 12/12.6a7672ab(the same code before the lean move): structured output 18/18, seeded sampling 9/9, behaviour 30/30, JSON 30/30, soak 60/60."draft": false.--prompt-cache-gib 0.--visionwithout the tower,--contextpast the caches, and a malformedSTEP5_LOAD_LIMIT_GBeach exit 1 with their message.step5_ref.py's top 5 next tokens in order, with logprobs within 0.02 (the sparse indexer choosing among 26k regions).zig build test test-goldenpass.tools/zig/lean_check.pyreports 517 problems, the same as 1.0.12.