Repository navigation
Ling-3.1-flash on CUDA: a shared KDA kernel first, then the family at TP=1 and TP=2 (EXL3) #850
Description
Ling-3.1-flash (inclusionAI/Ling-3.1-flash, MIT, released this week) is a 572B MoE: 72 layers, 63 of them KDA linear attention and 9 MLA (every 8th layer), the first 4 dense, then 512 experts (top 8, sigmoid, group-limited routing over 8 groups keeping 4) plus one shared expert, and one MTP layer. The release is FP8 (572 GB), served by SGLang and vLLM on 8x H200 or 4x GB300.
On Sparks it fits as EXL3: about 90 GB of weights at 1.25 bpw on one Spark, and 180 to 215 GB at 2.5 to 3 bpw on two. TensorFold serves across Sparks, so it is the engine where the two-Spark size runs. We'd like to build it, and would welcome help with any of the parts below.
Where we are (exllamav3, for the quants and as a reference). We're adding the architecture to exllamav3 so the packs can be made, checked against the bundled modeling_bailing_moe_v3.py (its own fla kernels):
- Layer 0 (KDA + dense MLP), 512 real tokens: within 0.36 % (relative L2) of the reference on the prompt path, and 0.35 % over 32 single-token decode steps with carried state (fp16 against the reference's bf16). Turning the safe gate off gives 16 %; swapping the forget and output gates gives 86 %.
- The EXL3 converter runs on it. Layer 0's state error at 1.63 / 2.02 / 4.02 bpw: rfn 0.138 / 0.089 / 0.021.
- MoE and MLA layers are next. The MLA layers add a per-head sigmoid gate on the attention output, which we've added to exllamav3's MLA.
What TensorFold needs on CUDA. On main, Nemotron is the only CUDA family. None of these exist there yet:
- A shared KDA kernel (prompt rows and decode rows) with verify windows that replay their rows exactly, so drafted equals serial. GLM-5.3 Flash and Kimi have KDA on Metal only, so this would serve them on CUDA as well. We'd take this first.
- MLA with the head-wise output gate, rope on 64 of the 192 query dims.
- Group-limited sigmoid top-k routing with an expert bias, next to the routing in cuda: sliding-window attention, sigmoid top-k routing and combine, dense rows for every family, checked against float64 #798.
- EXL3 projections at TP=1 and TP=2. @BHCC2025's GLM work in GLM-5.3-Flash and full GLM-5.3 on the native CUDA engine (tracker) #593 / glm: GLM-5.3-Flash (NVIDIA's NVFP4 checkpoint) as a CUDA family at TP over hosts #804 has both; Ling would use the same pieces.
- The family itself: config, weight map and layer graph, per
docs/recipes/adding-a-zig-family.md. MTP from its one draft layer comes later.
Where help is welcome: any of 2 to 4 as its own PR, and anyone with two or more Sparks willing to run the TP=2 checks. We have one Spark, so we can check TP=1 and the 1.25 bpw size ourselves, but not TP=2. We'll post the EXL3 packs once they check out against the reference: 1.25 bpw first, then about 2.5 bpw.
Question for @ashhart: is KDA as a shared CUDA module (item 1) the right first piece, or should it stay inside the family until a second CUDA family needs it?
Activity
jschmied commentedon Oct 15, 2026
Progress on item 1, and three questions before it becomes a PR.
KDA kernel (CUDA), branch
jschmied:cuda-kda, one commit onc7137bb3. It's a shared module next togdn_prep. It takes a family's projected rows and does the conv window, the safe or softplus gate, beta, the q/k L2 norms, the delta rule and the gated output norm.kda.cu: decode rows and verify windows. A window logs each row's decay, normalized key and update. The next call replays the kept rows from that log, so a window's rows have the same bits as those rows decoded one at a time.kda_prompt.cu: prompt rows in 16-row chunks from a segment's first row. A prompt cut at multiples of 16 rows has the same bits as one uncut call.core/kda_ref.zig: a float64 reference with a per-element bound.Checks (
tf-cuda-test kda,kda-real, GB10):chunk_kdatakes 14 ms for the same 2048 rows. No tensor cores yet.The exllamav3 side. All 72 layers now match the reference within 0.18 % (relative L2, routing-flip tokens excluded). That includes the MLA head gate and the clamped SwiGLU on layers 64 to 71: the bundled HF modeling skips the clamp, vLLM and SGLang apply it, and we follow them. Full-model logits on 512 tokens: top-1 agreement 99.2 %, KL 1.6e-4. The 1.25 bpw conversion is running (about 2.5 h per MoE layer) and should finish around Oct 22.
Questions for @ashhart:
zig/src/cuda/kda.zig) or live inside the Ling family until a second CUDA family needs it?jschmied commentedon Oct 15, 2026
A correction to my plan above: #804 (@Bizuayeu, GLM-5.3-Flash on CUDA) already has a KDA kernel, MLA on the 512-wide latent, sigmoid top-k routing with bias, and the clamped SwiGLU. Ling would extend those rather than add its own: rope dims and 128-wide heads for the MLA, plus the head gate and group-limited routing. I've asked on #804 whether that suits them. That makes question 1 above "our KDA or #804's, and where should the one we keep live?"
ashhart commentedon Oct 15, 2026
Hello @jschmied,
Thank you for the Ling-3.1-flash notes and the layer 0 check.
On main the CUDA server registers Nemotron only (native/cuda.zig:13). GLM and Kimi KDA there are Metal kernels (kda_rows.metal, kimi/kda.metal).
A shared CUDA KDA is the right first piece. #804 already runs GLM's (cuda_kda.zig:54), and Ling is the second CUDA family that needs it, so that kernel becomes the shared module, with verify windows so a drafted window matches the serial one. The KDA in #855 (zig/build/cuda.zig:51) is the same job and folds into that module.
MLA with the head-wise output gate, and the group-limited sigmoid top-k beside #798, can follow as their own pieces. EXL3 projections at TP=1 and TP=2 use the one EXL3 module (#832) and the one transport. The family, and its one MTP layer, wait on those. Your GB10 covers TP=1.
Thank you again.
jschmied commentedon Oct 15, 2026
Thank you, @ashhart, that settles the order. The Ling pieces are up as draft #874, stacked on #804 (GLM's bytes unchanged; @Bizuayeu agreed to this route on #804):
Its recipe,
docs/recipes/ling-3.1.md, lists what is still open as items anyone can take: lanes hooks and serving, MTP, TP through the one transport, an NVFP4 checkpoint, and runs on 3 to 4 Sparks. NVFP4 at about 324 GB is the size most people can run. One question is on #874: should group-limited top-k stay in #804's routing or move beside #798?ashhart commentedon Oct 16, 2026
Thank you, @jschmied, for #874 and for separating the KDA, MLA and routing work.
We rechecked e29e0c2. The zero-limit regression is closed and the WY experiment is removed. The branch now has lane serving and MTP, so the older missing-serving list is behind the code. KDA geometry remains a blocker at ling/config.zig:388, and the ownership, capacity and slot-initialization cases are listed on the PR.
Please keep the focused repair on #874, including rejection of unsupported KDA head widths and tap counts before launch, so we can recheck the same implementation. We'll come back on the shared routing module's placement.
Posted by Ash's AI Assistant Synthena
ashhart commentedon Oct 19, 2026
Thanks @jschmied for separating these pieces. We'll use one shared CUDA KDA module, extracted from GLM's existing chain and replay, with both families calling it. Group-limited selection belongs in the shared MoE routing module beside #798, with sigmoid-before-bias kept distinct from its existing logit-before-bias policy.
The focused extraction should preserve GLM's rounding, reduction order, replay and ungrouped routing bytes, add Ling's group mask, reject unsupported geometry before launch, and include serial/window, cut-prompt and routing comparisons. TF32 or reduced-precision prompt arithmetic needs a separate opt-in and quality gate.
Your pause on #874 stands; this placement decision doesn't ask you to resume it or reopen the whole-model work.
Posted by Ash's AI Assistant Synthena