Repository navigation
Conversation
40b39b4 to
191d8b7
Compare
|
Thanks @jschmied, this is a careful piece of work. The vision tower is checked kernel by kernel against float64, each Reading it on current main, the code holds up:
What keeps it a draft is the stack. It needs #1696 and #1502 first, and two of its first commits are now on main in |
Without torch (or with --numpy) the forward runs on numpy: MLX affine and bf16 weights are dequantized exactly; NVFP4 and FP8 stay torch-only and are refused by name. --bf16-points rounds where the Intel GPU engine stores bf16 (embedding rows and RoPE cos/sin too, attention P and logits never). dump writes the float64 hidden state after every layer and check replays an engine's greedy replies teacher-forced. The collect and miss-table CLI is unchanged.
…, gated GQA attention over a ring, sandwich norm, SwiGLU) and their FP64 op tests
…state and KV rings) loading by names.zig, its checkpoint test against ref64.py's float64 dumps
…pt projections on nem_prow_gen12 without DPAS, Iris Xe measurements
…pied on the device LFM3's per-stream state (the conv rows of 13 layers, the KV slots of 11, the position) moves out of the layers into a sequence the model binds; the engine makes, frees and binds sequences and saves and restores prompt states (the conv rows, the full layers' slots up to the mark, the sliding rings whole). The lane host serves LFM3 on one sequence until the engine verifies windows (a later commit).
… and RoPE with the KV store lfm3_mv.cl takes attn.cl's qmv_row verbatim (same compile options, so the same contraction) and runs w1 and w3 of a row with their SwiGLU in one launch, and q, k, v and the gate in one; lfm3_rope_store normalises and rotates the q and k heads and stores k and v in one launch, through the same per-head function lfm3_qk_norm_rope now calls. Both RoPE kernels require sub-groups of 16: their norm sums one partial a sub-group into part[4], which a compiler-chosen width of 8 would overrun and 32 would sum in another order. Every value is rounded as before: the greedy tokens of the three reference prompts are unchanged. Decode on Iris Xe: 15.2 to 17.2 tok/s (LFM3_FUSE selects the fusions for A/B timing).
On Iris Xe a prompt window of 128 rows beats 32 and 64 with the 8 x 4 prompt tiles (1,024-token prefill 84 against 78 tok/s at 64). The greedy tokens of the three reference prompts are unchanged at windows of 128 and of 64.
…ream, verified windows, shared forwards LFM3's engine takes the contract Gemma 4 E4B's takes on the XPU lane host (copy_only): one sequence a stream in up to 8 slots, the lane core's context copies verified in windows of up to 16 rows (verifyRows, keepRows), several streams' windows in one forward (verifyMulti, slotOf), window costs timed at load (1, 2, 4, 8 and 16 rows, the widths between interpolated: 8 s to load instead of 21 s, the same round plans), and prompt states kept anywhere (cache_rules). One line in the adapter (xpu_nem.zig): the host takes a family's copy_match when its table row declares one, before XPU_COPY_MATCH. LFM3's is 8: at the core's 4, copies cost a short chat 6% on Iris Xe; at 8, prose and code decode 65% and 38% faster than with --no-drafts, and the two chats within 2% of it (session noise). A window's rows sum as the decode steps do: its projections and head run on pdense, whose bf16 and fp32 outputs equal qmv4_bf and qmv4_f32 bit for bit (tf-xpu-lfm3_rows-test). Conv and attention run on each stream's own state and position; lfm3_conv_hist also records each row's B * x, so keeping the first c rows of a window restores the conv state those rows leave. Decode ids are kernel arguments, so no host memory the caller may reuse is read later. tf-xpu-lfm3_streams-test checks streams in shared forwards, freed and joined mid-way, and windows rolled back to their kept rows, against each stream alone.
…les, with copy drafts, streams and the prompt cache
xpu_model.zig drops its own invFreq (theta^e in f32, 1-3 ulp from torch in 16 of the full layers' 32 frequencies and 28 of the sliding layers' 64) for Rotary.inverseFrequencies() (f64 pow, then f32: torch's values), and sizes the sliding rings and windows by cfg.ring and LayerType.window, as the CUDA backend does.
…loader (2.7x TTFT at 1k), verified windows keep pdense pgemm (lfm3_mv.cl) tiles 64 x 64 in SLM as vit_gemm does and dequantizes q * scale + bias as it loads a weight tile; each output sums an fma chain a 64-wide k block and adds the blocks in order, so a row's sum does not depend on how the prompt is cut. Model.forwardPrompt routes every projection of a prompt window to it (prefill, prefillSpan, prefillWindows); verified windows, shared rounds and decode keep pdense and qmv, so drafts stay bit-equal to one-row steps. lfm3_pgemm checks it against FP64 in groups of 64 and 32 (the other group size as the planted control) and every row alone = in a window; lfm3_chunks checks the model's windows of 64, 128, 256 and a ragged cut bit for bit (rows, logits, the next step), pdense's order as the control. On Iris Xe through the server, TTFT at 1,024 tokens goes from 12.3-12.4 s to 4.6-4.7 s and at 8,192 from 116.3-116.7 s to 55.4-55.6 s.
|
Thanks @ashhart. Rebased as @Mirrdhyn, you run XPU on an Arc Pro B70. Would you have time for a second Intel GPU's run of this PR? It adds the LFM3-VL vision tower and LFM3's text decoder to the XPU backend; so far it has only run on an Iris Xe (Gen12LP, no DPAS), and a Xe2 run would show both that it works with DPAS and that nothing existing moves. On branch
Just the outputs pasted here would be enough. Thank you! |
|
Thanks @jschmied for the quick rebase. The three commits that landed as #1839 and #1692 are gone, and the series now applies cleanly on current main with nothing else changed. Thanks too for asking @Mirrdhyn for an Arc run; those three checks are exactly the ones we'd want. So what's left is outside your commits: #1696 landing first (and #1502 for the two carried files), plus the Iris Xe cells for this head when they're ready. Once those are in, this can leave draft. |
…apper, as main's lfm3 config now imports rope_freq
…embedding, group 64 unchanged The matvec sources (attn.cl's qmv_row, lfm3_mv.cl, nem_prow.cl's pdense) take the MLX group size as QG, 64 by default; attn_g32, lfm3_mv_g32, nem_prow_g32 and nem_prow_gen12_g32 are the same sources with QG=32. Every existing SPIR-V image is byte-equal to before, so LFM3-1.7B's forward and tokens do not move. lfm3_q8.cl reads the 8-bit embedding in groups of 32: embed8, embed8_tok and the tied head qmv8_f32; a verified window's rows take qmv8_f32 one row at a time, as the decode step sums them. The engine admits LFM3-1.7B (4-bit, group 64) and LFM3-VL-2B's text decoder (4-bit, group 32, the embedding 8-bit). tf-xpu-lfm3_q32-test checks the group-32 kernels and the 8-bit ones against float64, pdense against qmv4_bf bit for bit, and fails the planted controls (groups of 64, the bytes of a word reversed). On the VL checkpoint, tf-xpu-lfm3_fwd-test against ref64.py's float64 dumps gives float64's next token on capitals, chat and items, every layer at 0.94-1.05x the bf16 gap.
zig/src/families/lfm3/image.zig and its test root zig/src/lfm3_image_tests.zig byte-identical to ashhart#1502 at f066644, so they drop out when ashhart#1502 lands.
…ild file of ours zig/build/lfm3_xpu_host.zig declares the step only while no other build file has, so ashhart#1502's own lfm3_host.zig step takes over when it lands; build/sys.zig's walk gives the test root its sys import.
…ted controls lfm3_vit.cl rounds to bf16 where transformers' bf16 tower does and computes in fp32 between: - vit_gemm: y = bf16(x . w^T + bias), 64 x 64 tiles in local memory, 4 x 4 a work-item (the patchifier, every projection and the projector's two layers); - vit_rms (bf16(w * bf16(x * rsqrt))), vit_add_rms (the residual step and the next norm), vit_layernorm with bias, vit_gelu with erf, vit_swiglu; - vit_dwconv: bf16(C * bf16(dwconv3x3(bf16(B * x)))) over the 2 x 2 merge-order rows; - vit_qk_norm_rope: the q/k head RMSNorm at head dim 64 and axial RoPE (rows, columns); - vit_attn: unmasked GQA over one picture's patches, online softmax by 16 keys, no score matrix. tf-xpu-lfm3_vit-test checks each against float64 on random data; every planted control fails: the table transposed, a centred (or uncentred) norm, the norm before the residual add, silu for GELU, gate and up swapped, the conv grid fed in raster order, rows and columns swapped in RoPE, a causal mask, and query head h on KV head h % 8.
…ainst ashhart#1502's float64 reference xpu_vision.zig loads the 138 bf16 tensors by ashhart#1502's names and shapes and runs one picture's merge-order patches to its gh * gw / 4 image rows: patchifier and embedding norm, the 14 layers (gated depthwise conv or GQA attention with QK norm and axial RoPE, each with the SwiGLU FFN), then the projector (LayerNorm a patch, 2 x 2 blocks as 4096-wide rows, linear_1, GELU, linear_2). zig/src/lfm3_image.zig roots ashhart#1502's image.zig where its relative imports reach. tf-xpu-lfm3_tower-test writes the patches, every layer's rows and the image rows, and the same patches with the grid passed transposed. tools/lfm3/vision_check.py runs ashhart#1502's numpy reference (vision_fp64_reference.py, unchanged) with its bf16 points and in plain float64: - 14 x 20 synthetic picture: image rows rel 4.21e-2, cos 0.99912 against float64 (the reference's bf16 points: rel 6.55e-2); the transposed grid: rel 0.82, cos 0.67; - a 3000 x 4000 photo of a form at 1,024 image tokens (72 x 54 patches): rel 2.84e-2, cos 0.99960 (bf16 points: rel 3.13e-2); transposed: rel 1.09, cos 0.41. The tower takes 164 ms on the synthetic picture and 2.97 s on the photo on Iris Xe.
…ed in place of the image tokens xpu-native reads lfm3_vl checkpoints with the LFM3 family (vl_model_type) and offers image input as request unless --vision-image-tokens names 64 to 4096; LFM3-1.7B still refuses --vision. The engine loads the tower with scratch for that budget. Before a stream's prompt windows, the adapter hands its pictures to the engine (pictures, when the engine has it): the tower writes their rows, and each prompt window copies the rows over the embeddings of the image tokens it holds. Decode, copy drafts, streams and the prompt cache are unchanged; a kept state after a picture is keyed by the lane host's image hash. The server's core module owns core/picture, so ashhart#1502's image.zig, which reaches it by relative path, is compiled from a copy the build writes with those two imports read as core's; the file itself stays byte-identical to ashhart#1502's. tf-xpu-lfm3_vl-test runs one picture without the server (tower, projector, text prefill, decode). Through the server on Iris Xe, greedy: a vaccination record photo, a receipt and an A4 page give the same replies twice (the second from the prompt cache) and with --parallel 2 all at once; a receipt after another picture with the same text equals it on a cold server. LFM3-1.7B's gates do not move: B2's tokens, --parallel 1/2/4 against the pinned solo replies, the cache resume.
…s within 1.7e-3 of float64 Every activation of the tower and the projector stays fp32 (weights bf16 as stored, sums fp32 as before); only the image rows go to the decoder in bf16, as ashhart#1502's rows buffer holds them. Each kernel keeps a rounding flag for transformers' bf16 points (vision.Round, Tower.round; 0 by default): with every bit set the forward is C11's bit for bit. The ablation on ashhart#1502's 14 x 20 fixture against vision_fp64_reference.py in plain float64: no one rounding point carries the gap (any one alone in fp32 leaves the rows at 2.7e-2 to 7.0e-2); taken together they do: the linear outputs, the conv's inner points and the attention path to 7.5e-3, the projector's three points to 3.3e-3, the pixels to 1.65e-3 (layer 13 at 4.3e-6). - fixture: image rows rel 1.65e-3, cos 0.999999 (C11: 4.21e-2, 0.99912); transposed grid 0.82; - the 3000 x 4000 form photo at 1,024 tokens: rel 1.66e-3, cos 0.999999; transposed 1.09. tf-xpu-lfm3_vit-test now runs the kernels with fp32 in and out, within 1e-5 of float64 (at most 1.6e-6), and its controls still fail, GELU's tanh form among them. tf-xpu-lfm3_tower-test dumps fp32 and the first layers' stages by ashhart#1502's trace names; LFM3_VIT_ROUND sets the rounding points.
…ed row for ashhart#1502's compare_logits.py tf-xpu-lfm3_vl_logits-test runs ashhart#1502's real-image fixture (90 tokens, 70 image rows at 4, the suffix "Describe this image." -> "The image contains a cat.") through the tower and the decoder in windows of 16 rows, and writes each row's logits as bf16, the projected rows and the patches. On a stand-in for ashhart#1502's cat (placecats.com/300/200 serves another JPEG now: 14 x 20 patches, 70 rows), against ashhart#1502's fp64_reference.py, compare_logits.py and vision_fp64_reference.py, unchanged: - the XPU's own rows into float64 (decoder alone): rel 1.72e-2 over 90 rows (ashhart#1502: 1.30e-2), 3 top-1 differences, 1 disjoint top set, all at image positions; the suffix: rel 1.17e-2, none; - float64's own vision (whole pipeline): rel 3.74e-2, 3 disjoint top sets over 90 rows (ashhart#1502: 7), all at image positions; the suffix: rel 1.30e-2, no top-1 difference (ashhart#1502: 2 disjoint).
…hes file An optional fifth argument, fp32 [patches][1536] for the picture's grid, replaces the decoded pixels: a receipt's own patches go through the tower and the decoder unchanged, e.g. ashhart#1502's real-image fixture made with the pixel normalization before 50b4e0d (patches sha f1c55a0dd28b as its .npy).
…row layout for every backend vision_spec.zig imports only std and config.zig: vision_config's fields (read from config.json), the 138 vision tower and projector tensors' names, shapes and dtype built from them, and the image rows' Span arithmetic (overlap, place). vision_spec_test.zig checks them against config_test's VL fixture, a real export's model.safetensors.index.json when TF_LFM3_VL_DIR names one, and planted renames.
…ws with it xpu_vision.zig takes its widths from vision_spec's lfm3_vl_2b and loads each tensor by spec.tensor (name, shape and bf16 bytes; the attention layers by spec mixers, 138 by spec.count), without a prefix buffer of its own. Engine.pictures packs a prompt's pictures with spec.place (the over-budget and too-many checks included), and the prompt windows copy image rows by spec.overlap. The tower's rows, the image-context logits and the served replies are bit-identical to before.
…olled full-tile P.V, opt-in Lane l holds dims 4l to 4l + 3 of each row's q, shuffled to every lane for its scores: no SLM, no spill on tgllp. Offsets are 32-bit from the kernel arguments; a full tile's 16 keys unroll so each shuffle is a register region. Every row's ops and order stay vit_attn's (bit-equal by construction).
…it_attn Each against float64 and planted controls (a causal mask, KV head h % 8, and the new one: q's rows read transposed), and each candidate's bytes against vit_attn's.
… 1.7x on Iris Xe The tower launches vit_attn_qr (q in registers, no SLM, no spill on tgllp's IGC 1.0.15468) at local size 128 in place of vit_attn; grid and arguments are unchanged. Each row's ops and order are vit_attn's, and lfm3_vit checks the bytes on the device. On Iris Xe at 1,156 / 1,400 / 8,580 patches it runs at 360 / 372 / 313 GFLOPS against 207 / 209 / 183, and every tower layer's rows and the image rows of the receipt, the A4 page, the vaccination photo and the fixture are byte-equal.
|
The head moved to Iris Xe (Gen12LP, WSL2), on
Spark (GB10), |
|
Thanks @jschmied. The rebase and the Iris Xe cells answer our points. This still waits on #1696, and on #1502 for the two carried files, which now conflicts with main. After those land, please rebase. The one cell still open is Nemotron's tokens against main on an Intel GPU. We don't have one, so @Mirrdhyn's Arc run, or any Intel GPU that can hold Nemotron, would close it. |
|
@jschmied @ashhart Arc Pro B70 run of 1. Nothing existing moves
Decode 147–180 tok/s on both, device memory 19.19 GB on both. 2. The new kernels (
3. One picture through the server Every item and the total are read correctly. 377 prompt tokens, 69 generated, TTFT 0.26 s, 252 tok/s. The same request twice gives the same |
|
Thank you @Mirrdhyn, that's exactly the cell that was open. On the Arc Pro B70, Nemotron 3.5 Lightning gives the same @ashhart, that covers Nemotron against main on an Intel GPU. The |
|
@jschmied A heads-up on what we'd like to do on top of this once it lands, so we don't step on your rebase. On the B70 at Decode with What we'd take, Xe2 only, behind the device check so Gen12LP keeps your kernels:
Each against your float64 references and with your tests unchanged. We'll prototype locally and open nothing before #1838 is in. Does that split work for you, or is any of it already on your list? |
|
@Mirrdhyn, yes, that split works for us. All four are on Xe2 behind the device check, so Gen12LP keeps its kernels, and they're checked against the float64 references with our tests unchanged. None of them was on our list except the spill, and we found that one's cause:
|
|
@jschmied Where we are on the Xe2 side, so the rebase stays easy. All of it sits on a local branch on top of
Decode on DPAS by default on Xe2. The g32 decode now runs on DPAS GEMMs. A decode round reads the weights once for every stream: up to 32 rows, shared conv/rope/attention launches, and a shared 8-bit head. Prompts that arrive together, up to 16, share their prefill. Served, B70 in a PCIe 5.0 x16 slot, LFM3-VL-2B MLX 4-bit:
Bits. A stream gives the same bits alone or shared, and a prompt gives the same bits however it's cut. The next token still matches the float64 reference. Nemotron's Next: more than 32 streams, and register pressure in the decode GEMMs. |
|
Thanks @jschmied, and thanks @Mirrdhyn for the B70 run: the same Nemotron tokens on main and the head, and main's images byte-equal, close the last open cell. Nothing more is owed here. Once #1696 and #1502 are in, please rebase and it's ready to land. The Xe2 follow-ups as you've split them, behind the device check, sound good to us. |
Stacked on #1696 (LFM3-1.7B on XPU, with the family's shared rope frequencies, ring and window) and #1502 (LFM3 + LFM3-VL on Metal).
Shared vision spec. The tower loads through
families/lfm3/vision_spec.zig(733e5c6c5): one backend-neutral table of the 138 vision tensors, plus the image-row layout (Span,overlap,place), checked against the real checkpoint index atc823fade98e3. The XPU loader, the picture placement and the splice use it (40b39b4d2). The tower rows, the cat logits and the three served replies are bit-identical to before. The same spec is ready for #1502's Metal loader and the CUDA tower.zig/src/families/lfm3/image.zigandzig/src/lfm3_image_tests.zigare carried byte-identical to lfm3: add native Metal vision support above the text decoder #1502's headf066644fa, in their own commit, so they drop out when lfm3: add native Metal vision support above the text decoder #1502 lands.zig/build/lfm3_xpu_host.zigand steps aside once lfm3: add native Metal vision support above the text decoder #1502'slfm3_host.zigdeclares it.What it adds
The VL checkpoint's text decoder on XPU: MLX 4-bit in groups of 32 and the 8-bit embedding (head tied).
QG(64 by default). The_g32images are the same sources withQG=32;lfm3_q8.clholds the 8-bit embedding rows and head.The vision tower and projector without DPAS (
lfm3_vit.cl,xpu_vision.zig): fp32 activations, bf16 weights as stored, fp32 sums. The kernels:Only the image rows go to the decoder in bf16.
Serving pictures:
lfm3_vlcheckpoints get lfm3: add native Metal vision support above the text decoder #1502's image offer (its processor, part type "image", 256 image tokens unless--vision-image-tokens, up to 4096).Tests:
tf-xpu-lfm3_q32-testandtf-xpu-lfm3_vit-test: every kernel against float64, each with a planted control.tf-xpu-lfm3_tower-testandtools/lfm3/vision_check.py: per layer against lfm3: add native Metal vision support above the text decoder #1502'svision_fp64_reference.py.tf-xpu-lfm3_vl-test: one picture end to end.tf-xpu-lfm3_vl_logits-test: every teacher-forced row's logits, for lfm3: add native Metal vision support above the text decoder #1502'scompare_logits.py.Checks (Intel Iris Xe, Gen12LP, WSL2)
Kernels. The fp32 vision kernels are within 1e-5 of float64; the worst is 1.6e-6 (attention over 1,156 patches). Every planted control fails:
The group-32 and 8-bit kernels match float64, with
pdenseg32 equal toqmv4_bfbit for bit; their controls fail too (groups of 64, the bytes of a word reversed).Text decoder of the VL checkpoint against
ref64.py's float64: float64's next token on 3 prompts, every layer at 0.94–1.05× the bf16 gap.Tower and projector, image rows against float64. #1502's reference is run in plain float64; transformers' bf16 tower is estimated by the same reference's own bf16 points:
53153c22)Against #1502's own figures. #1502 measures native rows against the reference's bf16-points graph, its "idealFP64":
Logits with a picture in context: #1502's real-image fixture, the cat, 90 teacher-forced rows, its tools unchanged, its exact input (the receipt's patches
f1c55a0dd28b):9d39c115, ≤ 1.2e-7 per value from the receipt's), the figures are (a) 1.49e-2 and (b) 2.01e-2, again with 0 in the suffix.A receipt's first reply token. With the bf16 tower (C11) and with transformers' bf16, the receipt's reply starts with '1' ("1. Roggenbrot…").
Serving (
--vision --vision-image-tokens 4096 --context 8192, greedy, thinking off): a vaccination-record photo, a receipt and an A4 page.--parallel 2is byte-equal to solo.LFM3-1.7B (#1696): B2's 64 tokens on 3 prompts,
--parallel 1/2/4equal solo, cache resume equals cold, rows 28/28, streams ok.Lint: lean 0 new,
zig fmt,git diff --check.Numbers (Iris Xe, head
c96e5a4f4, one session, 3 runs each)Two speed changes since the first version, both measured against the old head in the same session (the laptop's GPU runs at a different clock from one session to the next: the old head's A4 tower took 11.0–11.3 s in one session and 15.3–15.5 s in another, so only same-session ratios compare):
Text prompts take xpu: LFM3-1.7B on Intel GPUs, Iris Xe first #1696's tiled prompt GEMM (
bb38b814b, a different summation order; its float64 gaps are in xpu: LFM3-1.7B on Intel GPUs, Iris Xe first #1696). Text prefill on the A4 page 2.5× faster.Tower attention takes
vit_attn_qr(c96e5a4f4): q in registers, 4 rows a sub-group, the full-tile P·V unrolled; no spill on tgllp, wherevit_attnspilled 896 B. Its output is byte-equal tovit_attn's on the GPU (tf-xpu-lfm3_vit-testat 1,156, 1,400 and 8,580 patches), and the tower's rows are byte-equal to before on all four pictures, every layer. 313–372 against 183–209 GFLOPS; the A4 tower 1.45×.Together: A4 TTFT 42.8 s → 22.8 s (1.9×) in one session. The served replies are unchanged by
c96e5a4f4.Peak device memory on the server: 3.14 GB with tower scratch for 4,096 image tokens and an 8,192-token context.
The fp32 activations cost 3–6% tower time against bf16 activations.