Skip to content

Zig CUDA ports: shared layers below lanes.Backend, and one place to coordinate them #548

Open

Description

@jschmied

Several model families are being ported to the Zig CUDA engine at the same time, and below the lane engine each port is building the same pieces again. The lanes engine, the CUDA runtime, the family registry, the checkpoint loader and the prompt cache are shared. Everything between the runtime and a family's Backend currently isn't.

What is being built more than once right now:

Piece Built separately in
Quantized linear layers (block-FP8, NVFP4, EXL3, MLX affine) #482 and our cuda-nvfp4 branch (FP8 and NVFP4 lane matmul), #496 (EXL3 decoder), the Flash Next port in #472 (MLX affine)
Grouped MoE experts and their routing plan our cuda-fp8-experts / cuda-nvfp4-experts branches (experts.route's plan + expert kernels), Nemotron's own experts, the Flash Next kernels on BobClawblaw's flashnext-zig-port branch
Copying the Python extensions' device code into the Zig build our generated copies, BobClawblaw's gen_copies.py (renaming the namespace or reordering instantiations changed two kernels' SASS)
Oracle fixtures and their test harness tf-cuda-test fixtures (#482), #507 (EXL3 fixture generator), the replay harness on flashnext-zig-port
Norms, RoPE variants, full and sliding attention, routers (softmax, sigmoid + bias), SwiGLU each family
Non-CUDA runtimes HIP (#463, #490), Level Zero (#522, #523)

The ports in flight: Flash Next (#472), Kolibri 1 (#481), GLM-5.3-Flash EXL3 (#509), the Qwen3.5 dense family (#503, #520, #528). The same direct_io test fix has also arrived three times (#486, #500, #540).

Proposal, in the order that would save the most duplicated work:

  1. One quantized-weight interface. A weight view per format, with one dispatch for decode rows and prompt rows, and each format's packing next to it. qmmf.zig in cuda: block-FP8 projections on the Python lane matmul's FP8G kernels #482's stack already does this for FP8G and NVFP4.
  2. One grouped-MoE layer. The routing plan plus per-format expert kernels behind one call. Our plan and FP8/NVFP4 expert kernels are byte-identical to the Python kernels on GB10 and could be the starting point.
  3. One kernel-copy tool and one fixture format. It should copy device code without changing its SASS, compare the SASS against the Python build, and write fixtures that every family's test command reads.
  4. Shared layer pieces (norms, RoPE, attention, routers), so a new family is mostly wiring.

Every shared piece would have to keep what the families have now: bit-identical output against the Python engine, and no slower than a specialised kernel. Where a family needs a specialised kernel, it can still register one behind the same interface.

Two questions for @ashhart:

  • Is this a direction you want? If so, which of the four first?
  • Would one tracking issue (this one, or a new one) work as the place where ports claim a piece before building it?

We can contribute 1 and 2 from what we already have, as small PRs. We'd hold them until you've set the direction.

Activity

  1. ashhart commented on Oct 10, 2026

    @ashhart
    Owner

    Thank you, @jschmied, for mapping where the ports duplicate each other. Yes, this is the direction: one model core that every family takes its GPU machinery from, with a family keeping only its own operators and weight handling, as CONTRIBUTING asks. Start with #558, the grouped expert plan as a shared module with Nemotron routed through it and its bits unchanged, then the one quantized-weight interface with decode-row and prompt-row dispatch per format. Each shared piece keeps the families' current contract: bit-identical output against today's kernels, and no slower than the specialised kernel it replaces, with a family still free to register its own kernel behind the interface. Please use this issue as the place where a port claims a piece before building it.

  2. Bizuayeu commented on Oct 10, 2026

    @Bizuayeu
    Contributor

    Claiming a piece here, as suggested: GLM-5.3-Flash on native CUDA from nvidia/GLM-5.3-Flash-NVFP4 (ModelOpt
    NVFP4, read as stored, W4A16), at TP=2 and TP=3 on GB10. It does not overlap the MLX affine port mentioned in #546 or
    the EXL3 port in #509: same family, different checkpoint format.

    Where it stands, on a branch over 1.0.2 (not yet public): the 45-layer forward runs on two GB10s over NCCL, one rank
    a process, each rank's partials gathered and summed in rank order. Against our Python 0.6.x engine on the same
    checkpoint, greedy decode gives the same token ids for three prompts of 512 tokens, with and without MTP drafts, and a
    prompt's last-row logits are byte-equal at 8k and 32k tokens. TP=3 runs on the Python engine and has not been moved
    yet.

    Two pieces in the table overlap ours, and we'd rather build on the shared ones than keep a second copy:

    What we'd send as small PRs, after the shared pieces above: the NVFP4 reader for the ModelOpt layout (the #308
    question, now on the Zig line) and the GLM-5.3-Flash NVFP4 family at --tp 2 and --tp 3, each gated on drafted
    equal to no-draft, resumed equal to fresh and concurrent equal to solo.

    Written with Weave (My Symbiotic AI)

  3. jschmied commented on Oct 10, 2026

    @jschmied
    ContributorAuthor

    Claiming the quantized-weight interface, next after #558 as you said.

    Plan, as small PRs in this order:

    1. cuda/qlinear.zig: one weight view and one dispatch. A tagged weight (affine4, fp8g, more formats later) with
      decode-row and prompt-row entry points. The first PR moves Nemotron's MLX affine-4 path (QLinear, the lane_gemv /
      split / cluster decode paths and prefill_matmul) behind it, with Nemotron routed through it and its bits unchanged.
      cuda: block-FP8 projections on the Python lane matmul's FP8G kernels #482's fp8.Lane goes behind the same interface unchanged. Receipts as in cuda: the expert plan as a shared module (cuda/grouped.zig); Nemotron routes through it #558: the tf-cuda-test fixtures, and
      Nemotron replies identical on main and the PR, greedy and sampled.
    2. NVFP4 behind the same interface: the lane matmul for decode rows and a prompt GEMM for prompt rows, each with an
      oracle fixture. It exists on our cuda-nvfp4 branch, SASS-identical to the Python build's kernels.
    3. Grouped experts per format (FP8, NVFP4) behind cuda: the expert plan as a shared module (cuda/grouped.zig); Nemotron routes through it #558's plan.

    @Bizuayeu, thanks. A GB10 run of cuda-nvfp4 and cuda-nvfp4-experts with nvidia/GLM-5.3-Flash-NVFP4 would be
    welcome; it would show early whether the weight view fits the ModelOpt layout as stored.

  4. Bizuayeu commented on Oct 10, 2026

    @Bizuayeu
    Contributor

    @jschmied Ran both branches on a GB10 (sm_121, nvcc 13.3, Zig 0.17.0): cuda-nvfp4 at bf8d716 and cuda-nvfp4-experts at cd4ee05. They branch separately from 36549fe, so each was built on its own. The Python reference was our fork's tree (release/2.8.0 lineage); your oracles ran there unmodified.

    Results

    Branch Test Fixture Result
    cuda-nvfp4 nvfp4 your oracle (synthetic) 9/9 PASS
    cuda-nvfp4 nvfp4 GLM layer 0 gate_proj [12288, 4096] 13/13 PASS
    cuda-nvfp4 nvfp4 GLM layer 0 down_proj [4096, 12288] 13/13 PASS
    cuda-nvfp4 nvfp4 GLM layer 3 expert 0 down_proj [4096, 2048] 13/13 PASS
    cuda-nvfp4 nvfp4 GLM layer 3 expert 0 gate_proj rows 0..704 (npad 768) 13/13 PASS
    cuda-nvfp4-experts nvfp4-experts your oracle (synthetic) 4/4 PASS
    cuda-nvfp4-experts nvfp4-experts GLM layer 3 experts 0..10, widths 2048 / 1024 / 704 5/5 PASS each

    GLM fixtures are nvidia/GLM-5.3-Flash-NVFP4 (snapshot 423acf37), lane and prompt at 1, 7, 16, 17, 40, 256 rows; experts at decode 1, 3, 40 rows and prompt 200 rows.

    • The tensors went into your format as stored: weight uint8 [N, K/2], weight_scale e4m3 [N, K/16], weight_scale_2 as the global scale. input_scale is not used. Only the MLP projections are NVFP4 in this checkpoint.
    • Cross-check against our own port (checked separately against the same Python engine): dense repack, lane and prompt bytes are equal, and so are the expert blocks and the decode outputs on routed slots.

    Prompt rows of the grouped experts: two forms

    For prompt chunks our GLM path does not use experts.cu. It uses a GLM-only experts_prompt.cu (ours, 3db60546 in our fork), so the 200-row prompt case on your branch differs from ours in 330 of 1,600 routed pairs at width 2048 (119 at 704). Swapping in experts_prompt.cu gives 0, so the kernel is the whole cause.

    experts.cu (your branch) experts_prompt.cu (our GLM path)
    Work unit items of 16 pairs items of 64 pairs x 128 columns, staged through shared memory
    Weights the same Experts4 blocks the same blocks, no second copy
    Prompt rows' bits the decode kernel's depend only on a pair's row and expert (chunk-invariant), not decode's bits, as with the MLX experts' prefill form

    Speed on the same GB10, layer 3's real experts, uniform random top-8 of 288 (not the real router), gate-up plus down, median of 30 runs (ms):

    Width Rows experts.cu experts_prompt.cu Ratio
    1024 (TP=2, served) 2,048 30.2 (13.7 TFLOPS) 10.9 (37.7 TFLOPS) 2.8x
    1024 512 10.7 9.2 1.2x
    1024 128 8.8 8.8 1.0x
    2048 2,048 63.6 19.9 3.2x
    704 (TP=3) 2,048 15.1 7.7 2.0x

    The staged form pays off on full prompt chunks and is even at 128 rows. Real routing is more uneven than uniform, so the 2,048-row ratio may move under real traffic.

    Within what we claimed above for GLM-5.3-Flash's NVFP4 family, we are ready to send experts_prompt.cu as a small PR behind your interface's prompt-row entry for grouped NVFP4 experts, once that entry exists, with an oracle fixture and the speed receipt. If you would rather keep one form, which one should the entry take?

    Paths your experts test does not reach, checked with local edits

    Three paths GLM uses are outside the committed test (it keeps Layer.limit 0, passes skip -1 and reads down_global directly). With a few local edits to the harness (optional limit / skip / plan-experts fixture params, a down_scale comparison, outputs filled with 0xFF before the launches; nothing committed to your branch), on layer 3's 288 experts:

    Path Setup Width 2048 Width 1024
    swiglu_limit 10 inputs scaled until about 16 % of gate and 33 % of up pre-activations clamp 6/6 PASS 6/6 PASS
    same inputs, limit 0 control 6/6 PASS 6/6 PASS
    limit 10 against limit-0 bytes proves the clamp fires fails as intended (5,475 of 9,216 words differ) fails as intended (2,703 of 4,608)
    skip a slot moe_block's 9-slot plan: top-8 plus the shared id 288, plan experts 289, skip 288 equal bytes; the skipped slot's rows stay untouched on both sides same
    down_scale your uploaded scales against make()'s bytes equal in every fixture equal

    Each PASS covers packing, decode at 1, 3 and 40 rows, and prompt at 200 rows (experts.cu on both sides). One harness limit showed up at width 2048: fixture.zig reads arrays with .limited(1 << 31), and 288 experts' packed_up is 2.7 GB, so the read stopped with StreamTooLong; raising the limit was the only other change.

    Logs, the fixture scripts, the harness diff and the benchmark are available if useful.

  5. jschmied commented on Oct 10, 2026

    @jschmied
    ContributorAuthor

    @Bizuayeu thank you, that is far more than a smoke test: real GLM tensors as stored, the cross-check against your own
    port, and the three paths our harness skips.

    On the prompt-row form for grouped NVFP4 experts: take yours. Nemotron's 4-bit experts already keep a separate prompt
    kernel (experts_prefill.cu, 64 pairs x 128 columns a CTA, next to the decode experts.cu), so a staged prompt form
    is the pattern main already has, and yours is 2.8-3.2x faster on full chunks. Our experts PR will keep experts.cu
    for decode rows only and leave the prompt-row entry to your experts_prompt.cu PR, with its oracle fixture and
    speed receipt.

    The 2 GB read limit in fixture.zig we'll raise in our NVFP4 PR. The first step of the interface is up as #587
    (Nemotron's 4-bit projections and #482's FP8 behind one weight view); NVFP4 goes behind it next.

  6. Bizuayeu commented on Oct 10, 2026

    @Bizuayeu
    Contributor

    @jschmied thank you, that settles it. Here is what we'll do on our side, in order:

    1. Nothing from us before the entry exists. We'll follow cuda: one quantized-weight interface (qlinear.zig); Nemotron's 4-bit projections route through it #587 and your NVFP4 PR, and rerun your NVFP4 branch on GB10 with GLM-5.3-Flash's real tensors when it's up (the same fixtures as last time, plus the limit / skip / down_scale cases), and post the results here.
    2. Then one small PR for the prompt-row entry of grouped NVFP4 experts:
      • experts_prompt.cu behind your interface's prompt-row call, experts.cu staying on decode rows as you said
      • an oracle fixture in your harness's form (synthetic, and GLM layer 3's experts at widths 2048 / 1024 / 704)
      • the speed receipt, measured as above (layer 3's real experts, gate-up plus down, median of 30 runs, both kernels on one build)
    3. The GLM family's NVFP4 linears and experts will move onto your interface once NVFP4 is behind it, and our own copies of those kernels go away. Your kernels already give the same bytes as our port on GLM's tensors (dense lane and prompt rows, experts on decode rows), and we'll check against our Python engine again when we switch.
    4. The GLM family itself stays in its own PRs within Zig CUDA ports: shared layers below lanes.Backend, and one place to coordinate them #548, on top of the shared pieces once they're in main.

    Where the family stands, for the coordination here: it now runs on main at f1aef9c, on native CUDA at TP=2 on two GB10s behind lanes.Backend. Rank 0 leads and the other ranks follow its steps, until #552's multi-rank shape lands. Against our Python engine on the same checkpoint, greedy and sampled replies give the same tokens. On the CONTRIBUTING contracts for one stream, drafted output equals "draft": false, including copy drafts from the core's proposer. After a cancel in the prompt or in decode, a discard after a rejected draft, an eos or stop string, or a max_tokens edge, the next request equals a fresh one. The prompt cache (Snapshots) is next.

  7. jschmied commented on Oct 10, 2026

    @jschmied
    ContributorAuthor

    @BlivionIaG, #612 and #613 and our CUDA pieces are converging on the same layer from two sides, and I'd like them to fit together before either hardens:

    A split that keeps one of each:

    1. core/quant owns the host side for every backend: formats, Config, TP cuts, the Uploader. We'd add the two formats CUDA serves today as format modules there: block-FP8 (e4m3 with 128 x 128 weight_scale_inv) and NVFP4 (ModelOpt: uint8 [N, K/2], e4m3 [N, K/16], weight_scale_2). Each comes with its exactness test in the form of quant/exact_test.zig.
    2. The device layout and kernels stay per backend. CUDA's lane tiles and HIP's tiles differ, so qlinear / experts would take a core/quant weight (its Bufs and format) and repack it on the device, as Affine4Kernels.pack does today for MLX. HIP's kernels slice would do the same with its own tiles.
    3. One rule for which kernel runs. feat(core): a kernel registry that picks a product's tile by tuned cost, one family of bits at every row count #613's family rule (a row's bits never depend on how many rows share a launch) is the contract our CUDA kernels already keep. Their K slices come from the weight's shape alone (qmm.split_k), so a row's bits are the same at every row count; ashhart checked FP8 from 1 to 300 rows on cuda: one quantized-weight interface (qlinear.zig); Nemotron's 4-bit projections route through it #587. Registering CUDA's entries in the same registry can come after; nothing in the split depends on it.

    If that's the direction, we'd send (1) as small PRs on top of #612 once it lands, and move qlinear / experts onto core/quant weights after that. Nemotron's bits are checked unchanged at each step, as with #587 and #605. Happy to adjust to whatever shape #491 settles on.

  8. jschmied commented on Oct 10, 2026

    @jschmied
    ContributorAuthor

    Claiming item 4 on CUDA, the shared layer pieces, driven by Kolibri 1 (#481) as their first user: RMSNorm and the fused add + RMSNorm, QK norms with RoPE and the cache write, the sigmoid router with bias and its top-k, the experts' weighted combine, then sliding-window and tree attention for decode and verify rows.

    Each one is our own CUDA kernel in cuda/, checked against a float64 or host reference inside its test (as in #633), with a row's bits the same at every row count, and drafted == plain on the family that uses it. They go up as small PRs, simplest first (the norms, combine and top-k). Nemotron's own glue stays as it is unless you'd like it moved onto them.

    This is CUDA only. #617 covers the same operations on HIP, and the shapes and contracts will match where they can.

  9. ashhart commented on Oct 10, 2026

    @ashhart
    Owner

    An update for coordination, @Bizuayeu and @jschmied. NVIDIA's NVFP4 checkpoint of Nemotron 3.5 Lightning now serves on our side, on #595's and #629's pieces, with Grouped.prompt running NVFP4 prompt rows through the decode form with bf16 sums. It's exact (drafted == plain, resumed == fresh, concurrent == solo), but on a GB10 that decode form takes 79% of an 8k prompt pass, so prompts run at about 0.45x of the MLX 4-bit path.

    To land it, we're now writing a staged NVFP4 expert prompt kernel for Nemotron's relu² experts, from our MLX staged kernel. If yours (experts_prompt.cu) is close, please say so here and we'll compare before either lands, so we don't end up with two. If it covers SwiGLU for GLM, both can live behind the same Grouped.prompt entry.

  10. ashhart commented on Oct 10, 2026

    @ashhart
    Owner

    NVFP4 Nemotron is in main (a6aa478). NVIDIA's nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4 now serves on the native CUDA engine, built on #595, #629 and #633. The staged NVFP4 expert prompt kernel we mentioned (relu² up, SwiGLU gate-up and bf16 sums, reading the decode form's blocks with no second weight copy) brought prompts on a GB10 from 0.45x to 0.99-1.00x of the MLX 4-bit path. A pair's bits depend only on its own row and expert, so chunking and batching can't change them. quant-ref now also checks NVFP4 SwiGLU prompt rows at 300 rows against float64.

    @Bizuayeu, if your experts_prompt.cu is faster on GLM's SwiGLU experts, we'd gladly compare it against this one on a real GLM NVFP4 checkpoint. The SwiGLU path here has so far passed only the synthetic float64 test.

  11. Bizuayeu commented on Oct 10, 2026

    @Bizuayeu
    Contributor

    @ashhart @jschmied, apologies: I missed your 20:54 note and opened #641 a few minutes before a6aa478 landed, which is exactly the two-kernels situation you wanted to avoid. I'm closing #641.

    Instead, we'll compare main's staged NVFP4 prompt kernel with ours on a GB10, on GLM-5.3-Flash's real NVFP4 experts:

    • layer 3, widths 704 / 1024 / 2048, with swiglu_limit, a skipped slot and down_scale;
    • chunks of 128, 512 and 2,048 rows.

    We'll post the speed here, along with whether its prompt rows match our 2.x engine's bytes. If ours turns out clearly faster on SwiGLU, we'll propose it as a change to the kernel in main rather than as a second one.

  12. Bizuayeu commented on Oct 10, 2026

    @Bizuayeu
    Contributor

    Results on a GB10 with GLM-5.3-Flash-NVFP4's real layer-3 experts (all 288, swiglu_limit 10, widths 704 / 1024 / 2048, chunks of 128 / 512 / 2,048 rows): main's staged kernel and ours run at the same speed, within 2% at every shape. So there's no case for a second kernel or for changing yours, and the GLM family will use main's.

  13. jschmied commented on Oct 11, 2026

    @jschmied
    ContributorAuthor

    A small follow-up to #633 that also starts the format modules I proposed for core/quant. The float64 decoders for FP8 and NVFP4 currently exist in three places: tests/cuda/quant_ref_tests.zig (e4m3, e2m1), tests/cuda/modelopt_tests.zig, and cli/cuda_precision.zig (which also has toE4m3).

    What I'd add:

    • core/quant/fp8.zig: e4m3, toE4m3, and reference(h, out []f64) for block-FP8 (128 x 128 weight_scale_inv).
    • core/quant/nvfp4.zig: e2m1 and reference(h, out []f64) for ModelOpt NVFP4 (uint8 [N, K/2], an e4m3 scale every 16 inputs, a global scale).

    Both mirror the existing mlx.reference and dense.reference. The three copies would call them, with a host test against hand-decoded codes. The CUDA packers in cuda/fp8.zig and cuda/nvfp4.zig stay where they are, since their layouts are CUDA's.

    This is host-only and backend-neutral, so a Metal or HIP port of those formats could check against the same references. I'll open it as one small PR unless you'd rather core/quant stay MLX/dense-only for now.

  14. jschmied commented on Oct 11, 2026

    @jschmied
    ContributorAuthor

    A proposal for the shared kernels: let every prompt scan write its recurrent state at a few requested rows, the way Flash Next's already does. A family could then keep a prompt-cache state anywhere without ending a chunk there.

    What exists

    • Flash Next does this today. pf_gdn_pre writes the conv window at a mark row (fn_prompt.metal L166-171). pf_gdn_scan/scan4/scan8 write the state to ST[m] when MK[m] == r + 1 (L246-248). That allows up to MARKS = 4 marks per call without ending the call (marks.zig L10-11).
    • Qwen 27B ends a chunk at each mark (qwen27_slots.zig L47-49), so a mark costs an extra prompt call: one more pass over the weights. warm hides most of that by keeping the state in a background pass after the reply. A request that lands during that pass still pays for it.
    • The other scans keep their state in registers and write it only at the end: tf_delta_gdn_chain_wide (deltanet.metal L284), GLM's glm_kda_scan, and CUDA's scan_rows.cu. Nemotron and GLM are planned, so their marks floor to chunk starts.
    • Nothing recurrent is shared yet apart from the cache itself. Each family has its own scan, its own conv-tail handling and its own Snapshots.

    What I'd propose

    One convention for prompt scans: a small mark list (MK[4], the row after which to store, 0 for unused) and per-mark outputs for the state and the conv window. This is the Flash Next layout, kept bit-for-bit what a chunk ending at that row would leave. The cost is one extra store per mark row, and nothing changes when no marks are set.

    Order, smallest first:

    1. core/deltanet.metal's prompt chain, the one shared DeltaNet library. Qwen 27B then keeps its marks inside a prompt call, as Flash Next does, and qwen27_slots stops cutting at them.
    2. GLM's KDA scan and Nemotron's scans (Metal SSD and CUDA scan_rows), if you want those families off chunk-start-only marks.
    3. A small shared helper for the snapshot copy (recurrent state, conv window, attention prefix), so each family only lists its layers.

    Each step would carry a test: the state written at a mark is byte-equal to the state left by a chunk ending there, and a prompt resumed from it equals a fresh one.

    I haven't measured the 27B side. I can't run the 27B here right now, so I'd need you or someone with a Mac to measure TTFT on a warm agent loop. Before writing step 1, does this convention suit you, or would you rather keep marks per family?

  15. 66 remaining items

  16. jschmied commented on Oct 30, 2026

    @jschmied
    ContributorAuthor

    /atc renew

  17. ashhart commented on Oct 30, 2026

    @ashhart
    Owner

    Status: step 1 of the Nemotron move has landed in #1353. The MoE residual's norm, Mamba's group norm, the MTP input
    norms and the bf16 token rows run on the shared kernels, in Nemotron's summation order, so replies are the same
    tokens and one fewer CUDA image loads.

    Step 2 (#1354, the router top-k) has passed its GB10 runs (tokens equal, decode level) and lands tonight. Left:
    step 3, attention, and the other families' moves. Thanks @jschmied.

  18. ashhart commented on Oct 31, 2026

    @ashhart
    Owner

    Status: step 2 of the Nemotron move has landed in #1354. Nemotron's router top-k now runs on the shared MoE top-k,
    and Bolt's latent-expert sum on the shared combine. Picks, weights and replies are the same bytes as before. In a
    GB10 trace of Lightning decoding, the top-k takes 2.74 µs a call instead of 4.37.

    Left: step 3, attention (with the comparison against block and window attention first, as you planned), and the
    other families' moves. Thanks @jschmied.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions