[jschmied]

Part of #1488. Stacked on #1377 (BlivionIaG's LFM2 stack, head 70f74ca34): the 10 commits below 7d7d8f49e are #1377's, and this PR adds only the last four. Draft until #1377 lands, then I'll rebase onto main.

What

LFM2.5-24B-A2B (lfm2_moe) served natively on CUDA from LiquidAI/LFM2.5-24B-A2B-MLX-4bit: plain serving with --parallel N and TF's built-in drafts. No DSpark drafter yet. The engine reuses the shared CUDA pieces wherever they fit.

• Reused as they are (13): the affine embedding rows, the 4/6/8-bit affine linears (prompt GEMM and GEMV), addRmsTail norms, the SwiGLU act with group sums, qk_rows (q/k norm + RoPE), the page KV write, the Affine4 experts with the grouped plan, the MoE combine, and tree_round argmax / sampling.
• New instances of shared templates (2): prefill_attention's prompt kernel at head dim 64, and the paged shared-window attention at head dim 64. Its body moves from Nemotron's file into a neutral paged_attn_segs.cuh. Head 64 exports tf_paged_attn_chunk_segs_64 / tf_paged_attn_merge_segs_64, and Nemotron's head-128 symbols keep their names and SASS.
• New kernels (3), under the kernel contract: the short conv with its rolling state, the conv-state commit, and the router. The conv and its commit live in a neutral liquid_conv.cuh, templated on the tap count (tf_liquid_conv_3 / tf_liquid_commit_3), so LFM3's 7-tap conv (Add native Metal support for LFM3-1.7B and LFM3-VL-2B MLX 4/6/8-bit checkpoints #1491) can reuse them; only the router stays in lfm2_moe.cu. The move is SASS-identical by function body on all five SMs, and served replies are byte-equal before and after. The router has fp32 weights, expert bias, normalized top-4 and eps 1e-6.

Commits:

• 7d7d8f49e cuda: the paged shared-window attention at head dim 64
• d70fc9821 lfm2_moe: CUDA glue kernels under the contract
• c6645063a lfm2_moe: serve LFM2.5-24B-A2B on CUDA
• 580d8f24b cuda: Liquid's short conv and its commit in liquid_conv.cuh, by tap count

22 files, +2,262 / −142 against #1377.

Verification (GB10, sm_121, nvcc 13.0)

• SASS identity: every existing function is unchanged on sm_80/86/89/120/121; 15 functions are added.
• Kernel contract: the 4 new kernels pass gates 1–4, and all 3 planted controls go red.
  • It found a real bug, now fixed: a row whose router scores were all NaN picked expert 0 twice.
  • Launch records: 30 of 30 pass on nvcc 13.0 and 13.2 (one 13.2 override).
• Served, real weights, --parallel 8, greedy:
  • drafted == plain on 8/8 prompts, including a partial acceptance;
  • concurrent == solo on 8/8, drafted and plain;
  • token hashes are identical across a restart.
• float64 reference (CPU, the MLX weights dequantized exactly, every op in float64): 61 of 64 greedy tokens agree.
  • Each of the 3 misses is the reference's second choice, with logit margins 0.009, 0.069 and 0.36.
  • A second reference rounds to bf16 at every point where the engine stores bf16, with the router, decode attention and the decode down output in fp32, as in the engine. Two variants of it serve as noise probes: every row on the decode path, and float32 arithmetic. At 2 of the 3 positions all three pick the engine's token. At the third (margin 0.36), both probes pick it, and the main variant keeps the float64 token by 0.30.
  • The rounding-only variants themselves disagree at 2 of 64 positions, so at these margins the misses are rounding on near-ties. I haven't compared against the engine's own logits: the CUDA backend doesn't return them yet.
• lean_check: 0 new problems against LFM2: shared-round decode #1377's head. zig build test passes.

Speed (GB10)

streams	decode tok/s (total)
1	87
4	194
8	303

Cold TTFT: 0.32 s at 2.8k prompt tokens, 1.32 s at 11.2k.

Not in this PR

• Pages leased from the shared KV pool: each stream has fixed pages for now. The prompt cache comes with it.
• The DSpark drafter on CUDA, after lfm2: keep a drafter's tapped rows on every forward #1390–lfm2: a cheaper DSpark hold (one concurrent encoder) #1394.

Questions

1. paged_attn_segs.cuh is still compiled into the nemotron_attention image. I'd move it and liquid_conv.cuh under zig/kernels/cuda/common/ once cuda: shared device primitives in zig/kernels/cuda/common (#1409, part 1: SASS unchanged) #1440 lands. Does that suit you, or would you rather each get its own image?
2. ATC shows one conflict with Nemotron: two members a pass at 8 rows a simdgroup for wide M1 windows #1466, in registry.zig (records), which is easy to resolve whichever lands second.

[ashhart]

Thanks @jschmied, this is a tidy port. It reuses the shared CUDA pieces wherever they fit, and the contract finding the all-NaN router row is a good catch.

#1377 has landed as rewritten commits, so this branch no longer merges with main: it conflicts in docs/kernel-contract.md, registry.zig and contract_limits.zig. Your own five commits cherry-pick onto main with one conflict, the image list at zig/tests/cuda/contract_limits.zig:28. With the first four on main, zig fmt, test-lfm2-moe-cuda-host, test-cuda-host and the Linux cross-compile pass.

Before this leaves draft:

1. Please rebase onto main.
2. The lean check has one new finding, the two-line comment at zig/kernels/cuda/nemotron_attention.cu:233. It is on the head too.
3. The three new attention launches, tf_paged_attn_chunk_segs_64, tf_paged_attn_merge_segs_64 and pattn_kernel<64, 8, 4, 8>, have limit records (lfm2_moe_cuda_records.zig:12-14) but no gate. The decode pair is new code: the SEG_WIDE 0 branch of paged_attn_segs.cuh:78-88. Please gate them against contract_attn.zig's float64 reference. Add a planted control that goes red, for example values read from the neighbouring lane's dims.
4. Please commit the float64 reference under tools/. Its receipt should list each top-1 miss with the prompt, the position, both tokens and both margins, and cover teacher-forced prompt rows as well as reply rows, as lfm3: serve LFM3-1.7B NVFP4 on CUDA, weight-only (stacked on #1492 and the LFM3 admission) #1508 does.
5. Docs:
   • a CUDA section in docs/recipes/lfm2.5-24b-a2b.md, with the command, the revision, the GB10 receipt and the weights' licence;
   • the README row (README.md:176 says one Mac);
   • the new checks in docs/kernel-contract.md.

To your questions: moving both headers under common/ after #1440 suits us. The #1466 conflict goes to whichever lands second.

Then we'll run the kernel contract and its limits at 13.0 and 13.2, and compare Nemotron's head-128 machine code with main. Then we'll serve the 4-bit export on a GB10: drafted against plain, concurrent against solo at 1, 4 and 8 streams, and speed.