[jschmied]

Part of #1488. Rebased onto main (80b30ed8a); #1377's commits are gone from the branch.

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, and its prelude into paged_attn.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 machine code.
• 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 (lfm3: serve LFM3-1.7B on CUDA: NVFP4 and FP8 weight-only, MLX 4-bit (stacked on #1513, #1492 and #1503) #1508) reuses them; only the router stays in lfm2_moe.cu.

Commits:

• 751ec79e9 cuda: the paged shared-window attention at head dim 64
• 3b613a26d lfm2_moe: CUDA glue kernels under the contract
• 5e5477f51 lfm2_moe: serve LFM2.5-24B-A2B on CUDA
• f3b5e60fb cuda: Liquid's short conv and its commit in liquid_conv.cuh, by tap count
• 22eba9470 kernel contract: the limits check reads lfm2_moe's fatbin
• e2fb1c492 kernel contract: the head-64 paged attention, decode and prompt rows, against float64
• f52c228a9 tools: LFM2.5-24B-A2B's float64 reference on the CPU, and its receipt of every top-1 miss
• 7390fb391 docs: LFM2.5-24B-A2B on CUDA, its receipt and the weights' licence

34 files, +3,171 / −251 against main.

Verification (GB10, sm_121)

• Machine code: against main, every existing function is unchanged on sm_80/86/89/120/121 under nvcc 13.0 and 13.2: nemotron_attention 45 functions (+10 new), prefill_attention 35 (+5), contract_controls 250 (+30).
• Kernel contract: the router, the conv and its commit, and the three attention launches pass gates 1–4. All 5 planted controls go red.
  • The head-64 attention is held head by head to contract_attn.zig's float64 bound (attention_ref), over shuffled 512-key pages with several streams a launch. Decode rows run the chunk and merge pair at 1/3,000, 16/1,500, 128/700 and 8/4,000 rows/keys. Prompt rows run pattn_kernel<64, 8, 4, 8> at 256 rows with no history, 64 rows after 3,000 keys, and 100 after 500. A NaN query, ±inf values and a NaN key are planted.
  • Its controls read values from the neighbouring lane's dims (decode, 15× the bound) and from the neighbouring 8 dims (prompt, 1,135× the bound).
  • It found a real bug, now fixed: a row whose router scores were all NaN picked expert 0 twice.
  • Launch records pass on nvcc 13.0 and 13.2.
• Served, real weights, --parallel 8, greedy:
  • drafted == plain on 8/8 prompts, including a partial acceptance;
  • concurrent == solo on 8/8 short and 8/8 long prompts;
  • token hashes are identical across a restart, and identical to the branch before the rebase.
• float64 reference (tools/lfm2_moe/ref64.py, CPU, the MLX weights dequantized exactly, every op in float64): the engine's greedy token on 421 of 435 teacher-forced prompt rows and 123 of 128 reply rows. With the engine's bf16 storage points it's 424/435 and 123/128. tools/lfm2_moe/REF64.md lists all 19 misses with both tokens and both margins.
  • The float64 margins run from 0.01 to 0.91. At about half of the misses the bf16-points pass picks the engine's token, so rounding alone decides them.
  • The four widest (0.63–0.91) may be routing flips under rounding, but I haven't checked that.
• lean_check: 0 new problems against main. zig fmt --check zig and zig build test pass.

Speed (GB10, before the rebase)

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.

Docs

• docs/recipes/lfm2.5-24b-a2b.md: a CUDA section with the command, the revision (667408d4), the GB10 receipt and the weights' licence (LFM Open License v1.0).
• README.md: the LFM2.5 row names one NVIDIA GPU with the 4-bit export.
• docs/kernel-contract.md: LFM2.5's kernels and the two attention controls under Coverage.

Not in this PR

• Pages leased from the shared KV pool, and the prompt cache: lfm2_moe CUDA: attention pages from the shared KV pool, and the prompt cache (stacked on #1492) #1503.
• 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.
• Moving paged_attn_segs.cuh 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.

[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 on CUDA: NVFP4 and FP8 weight-only, MLX 4-bit (stacked on #1513, #1492 and #1503) #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.

[jschmied]

Thanks @ashhart. All five are done, and the branch is rebased onto main (80b30ed8a); the head is 7390fb391.

1. Rebase: onto main, with the image list at contract_limits.zig resolved. zig/build/cuda.zig takes main's side, and lfm2_moe_cuda now registers in cuda_engines.zig.
2. Lean: the token-pasted names at nemotron_attention.cu:233 are literal macros now (SEG_CHUNK / SEG_MERGE), so kernel_symbols still finds both. lean_check has 0 new findings against main. Nemotron's machine code is unchanged (below).
3. Gates: zig/tests/cuda/contract_paged_attn.zig holds tf_paged_attn_chunk_segs_64 + tf_paged_attn_merge_segs_64 and pattn_kernel<64, 8, 4, 8> head by head to contract_attn.zig's float64 bound. They run over shuffled 512-key pages with several streams a launch, and a NaN query, ±inf values and a NaN key are planted. All three pass gates 1–4.
   • Two controls, built from the same bodies through a lane knob that is 0 in the shipped kernels: the decode reads its values from the neighbouring lane's dims (red at 15× the bound), and the prompt rows from the neighbouring 8 dims (red at 1,135×).
   • To build them from one body, prefill_attention.cu's helpers moved into prefill_attention.cuh, and the paged prelude into paged_attn.cuh. Every existing function is unchanged on sm_80/86/89/120/121 under nvcc 13.0 and 13.2: nemotron_attention 45 functions, prefill_attention 35, contract_controls 250, plus only the new ones.
4. Reference: tools/lfm2_moe/ref64.py with its receipt REF64.md. Float64 agrees on 421/435 teacher-forced prompt rows and 123/128 reply rows; with the engine's bf16 storage points it's 424/435 and 123/128. Each of the 19 misses is listed with case, row, position, both tokens and both margins.
5. Docs: the recipe has a CUDA section (command, revision, GB10 receipt, LFM Open License v1.0), README line 176 names one NVIDIA GPU with the 4-bit export, and docs/kernel-contract.md lists the new checks under Coverage.

The served checks were rerun on the rebased head (GB10, --parallel 8). drafted == plain 8/8, concurrent == solo 8/8 short and 8/8 long, and every token hash is equal to the pre-rebase branch's.