[orangezed]

📌 Description

FLARE_KIMI3_KDA_FUSED_VERIFY=1 (default off) is read when a verify layer with records is prepared, and must be set identically for export and serving. With it, each KDA verify layer runs b12x's one-launch verify (speculative short conv, recording verification and gated RMSNorm) instead of short_conv.run + kda_record. That is 3 launches -> 1 per KDA layer, over 69 layers. A new custom op flare::kda_record_fused mutates the output, the records and the convolution state. The fused binding is prepared by a finalizer after both bindings exist, and is re-prepared if they are replaced. Geometries the kernel refuses fall back to the three launches with a warning.

Kernel-level gains are in FlashInfer local-inference-lab/flashinfer#90 (rows 1: 6.90 -> 5.66 us; rows 8: 13.03 -> 11.28 us per layer).

Exactness: bit-identical to the three-launch path. Default-off guarantee: unset keeps the three launches and the existing op graph.

🧪 Tests

• New tests/recipes/test_kimi3_kda_fused_verify_cuda.py. It uses synthetic BF16 weights at the K3.1 TP8 rank geometry, compile_forward + CUDA graph and a poisoned arena, and compares off vs on bitwise for output projection, gated output, conv state, recurrent state and records at rows 1/4/8. GPU only; on this branch it was collected and skipped on CPU (actually run). The other kimi3 CPU tests have the same failure set as the release branch.
• GPU on the pre-rebase tree: kimi3 tests 30 passed, 16 skipped (checkpoint), and 1 failure (test_shard_columns_survive_export_scratch_clearing) that was present before the change, with the switch off and on.

Topology assumptions

Generic (single GPU; uses b12x thread-block clusters).

Reviewer Notes / risk

• With an older b12x that lacks Add a one-launch KDA speculative verify to b12x gdn_decode local-inference-lab/flashinfer#90 and the switch on, preparation raises AttributeError, because only ValueError is caught for the fallback. With the switch off nothing touches the new API. A hasattr guard could be added if the maintainer prefers.

[orangezed]

GPU validation of this branch as published (base release/k31-exl3-tp8-20261108 1d4007f), one RTX PRO 6000 Max-Q, with b12x from FlashInfer release/k31-exl3-tp8-20261108 plus the commit of the FlashInfer PR that adds gdn_decode.prepare_kda_fused_verify / run_kda_fused_verify. No changes were needed.

• FLARE_RUN_CUDA_TESTS=1 pytest tests/recipes/test_kimi3_kda_fused_verify_cuda.py: 3 passed (rows 1/4/8). The test builds compile_forward + CUDA graph over a poisoned arena and compares the switch off vs on bitwise for output projection, gated output, conv state, recurrent state and records.
• tests/recipes/test_kimi3_speculation_cuda.py and tests/recipes/test_kimi3.py with the Kimi K3.1 checkpoint mounted: the pass/fail set is identical with FLARE_KIMI3_KDA_FUSED_VERIFY unset and set to 1. One case fails with and without the switch, and it fails the same way on the release branch without this PR: test_checkpoint_kda_verification_commits_only_accepted_prefix, ValueError: tensor extent must divide into aligned rank partitions. That test uses the K3 test geometry against the K3.1 checkpoint. The checkpoint tests in test_kimi3.py fail the same way, again with and without the switch.
• Kernel-level exactness and timing are on the FlashInfer PR. On its rebased branch: bitwise identical to the three launches at rows 1..16, 2,000 graph replays per row count with 0 mismatches, and rows 1/4/8 at 6.89 -> 5.34, 9.17 -> 7.47 and 12.23 -> 10.49 us per layer.
• The AttributeError note in the description still applies: with the switch on and a b12x that lacks the FlashInfer PR, preparation raises. We did not change that.