[Tianlei Wu (tianleiwu)]

Summary

Implements Kimi Delta Attention (KDA, per-key-dimension decay linear attention) and supporting operators/optimizations for inclusionAI/Ling-3.0-tiny (BailingMoeV3ForCausalLM, bailing_hybrid) across CPU, CUDA, and WebGPU execution providers:

1. com.microsoft.GatedDeltaNet KDA Support (CPU, CUDA, WebGPU):
   • Supports rank-3 per-key-dimension decay ([total_tokens, num_heads_v, head_size_qk]), rank-2 dt_bias ([num_heads_v, head_size_qk]), and gate_activation="sigmoid" (alias "kda"): $g_{t,h,k} = \text{gate_lower_bound} \cdot \sigma(\exp(\text{a_log}h) \cdot (\text{decay}{t,h,k} + \text{dt_bias}_{h,k}))$ with $\epsilon = 10^{-6}$ Q/K L2 normalization.
   • CUDA EP:
     • Warp-specialized recurrent/decode kernel (GatedDeltaNetRecurrentKernel) with vectorized per-key-dim decay state update.
     • Two-stage split-chunked prefill engine (GatedDeltaNetKdaPrepareKernel + GatedDeltaNetKdaScanKernel, $B_T=16, D_K=128, D_V=128$) using exact forward-substitution $(I+M)^{-1}$, lower-triangular pair accumulation, 256-thread / DVB=16 tensor-core scan (1.40x faster KDA prefill on H200), multi-pass workspace carry, and ep.cuda.gated_delta_net_fp16_state_guard / ep.cuda.gated_delta_net_bf16_precision.
   • CPU & WebGPU EPs: Full KDA recurrent, varlen (cu_seqlens), and compact speculative state-replay capture support.
2. com.microsoft.LinearAttentionGates (CPU & CUDA):
   • Fuses low-rank bottleneck projection + per-key-dim decay & sigmoid beta gate computation.
3. com.microsoft.MoEGroupRouter (CPU & CUDA):
   • Fuses the 24-node noaux_tc group-limited sigmoid MoE router (Sigmoid -> Add(bias) -> TopK(group) -> ReduceSum -> TopK(groups) -> ScatterElements -> Mul -> TopK(k) -> Div -> Mul(routed_scaling_factor)) into a single warp-per-token kernel (MoEGroupRouterWarp128Kernel on CUDA, 9.0x faster decode router and -391 kernel launches/step across 17 MoE layers).
4. CUDA Kernel Optimizations for Hybrid Linear/Full-Attention MoE Models:
   • VarlenCausalConvWithState: Two-phase 2D token-chunk tiled prefill kernel (VarlenCausalConvBodyKernel + VarlenCausalConvPrefixAndStateKernel) achieving 39.8x faster prefill (26.29 ms -> 0.66 ms at 2k tokens).
   • PagedAttention INT8 KV Cache Decode: Per-warp register online-softmax + char4 vectorized cache reads in PagedDecodeSplitKV and parallel PagedDecodeReduce (2.44x faster INT8-KV decode).
   • QMoE Block-FP8: 8-wide vectorized QMoEFp8GemvKernel with hoisted block-scale index math (2.46x faster decode) and 16-wide vectorized QMoEDequantizeFp8WeightsVec16Kernel (7.06x faster prefill dequantization).

Validation

• Unit tests in onnxruntime_provider_test --gtest_filter="GatedDeltaNet*:*PagedAttention*:*QMoE*" (83/83 GatedDeltaNet tests passing, including GatedDeltaNetTest.MoEGroupRouterCpuAndCuda).
• End-to-end inclusionAI/Ling-3.0-tiny benchmarks & MMLU-Pro accuracy validated across BF16, Block-FP8, INT4+INT8KV, and INT8+INT8KV recipes against vLLM and llama.cpp.

[Copilot]

🟡 Changes recommended

Unresolved critical findings include an out-of-bounds CUDA cache read, WebGPU binding-limit and shader-validation failures, and decay calculations that produce incorrect results or NaNs.

7 open issues

• Propagate gate_lower_bound through split and BF16 paths · New
• Handle saturated per-key decay prefixes without NaNs · New
• Guard cache loads for channels beyond head_size · New
• Convert half-precision gates and specialize pipeline metadata · New
• Pack KDA gate bindings to support adapter limits · New
• Add explicit CPU tests for causal convolution with state · New
• Update operator schema docs and regenerate documentation · New

What changed in this PR

Adds Kimi Delta Attention support to ONNX Runtime’s GatedDeltaNet across CPU, WebGPU, and CUDA, including per-key-dimension decay and accelerated CUDA prefill.

Changes:

• Extends the operator schema, gate activation, and compact state-update layout for KDA.
• Adds CPU kernels, WebGPU gate support, CUDA split-chunked prefill, and KDA tests.
• Tunes CUDA block-FP8 MoE dispatch and paged-attention decode.

File	Description
onnxruntime/​test/​contrib_ops/​gated_delta_net_op_test.cc	Adds KDA reference, execution, and planner tests.
onnxruntime/​core/​graph/​contrib_ops/​bert_defs.cc	Extends gate attributes, bias shapes, and state-update inference.
onnxruntime/​contrib_ops/​webgpu/​bert/​gated_delta_net.wgsl.template	Implements per-key decay and sigmoid gates.
onnxruntime/​contrib_ops/​webgpu/​bert/​gated_delta_net.h	Adds KDA shader parameters and kernel fields.
onnxruntime/​contrib_ops/​webgpu/​bert/​gated_delta_net.cc	Updates validation, bindings, and dispatch for KDA.
onnxruntime/​contrib_ops/​cuda/​moe/​moe_quantization.cc	Limits fused FP8 GEMV to small decode workloads.
onnxruntime/​contrib_ops/​cuda/​bert/​paged_attention_impl.cu	Increases split occupancy and combines value-channel accumulation.
onnxruntime/​contrib_ops/​cuda/​bert/​gated_delta_net.h	Stores the configurable decay lower bound.
onnxruntime/​contrib_ops/​cuda/​bert/​gated_delta_net.cc	Validates KDA inputs and populates dispatch descriptors.
onnxruntime/​contrib_ops/​cuda/​bert/​gated_delta_net_split_impl.cu	Adds KDA prepare and scan kernels.
onnxruntime/​contrib_ops/​cuda/​bert/​gated_delta_net_plan.h	Selects the KDA split engine and sizes workspace.
onnxruntime/​contrib_ops/​cuda/​bert/​gated_delta_net_mma.cuh	Extends shared decay helpers and kernel parameters.
onnxruntime/​contrib_ops/​cuda/​bert/​gated_delta_net_impl.cu	Updates recurrent decay, capture, and tail dispatch.
onnxruntime/​contrib_ops/​cpu/​cpu_contrib_kernels.cc	Registers new CPU kernels and BF16 variants.
onnxruntime/​contrib_ops/​cpu/​bert/​varlen_causal_conv_with_state.h	Declares the CPU variable-length convolution kernel.
onnxruntime/​contrib_ops/​cpu/​bert/​varlen_causal_conv_with_state.cc	Implements stateful CPU convolution and compact capture.
onnxruntime/​contrib_ops/​cpu/​bert/​linear_attention_gates.cc	Adds BF16 gate and normalization registrations.
onnxruntime/​contrib_ops/​cpu/​bert/​gated_delta_net.h	Declares the CPU GatedDeltaNet kernel.
onnxruntime/​contrib_ops/​cpu/​bert/​gated_delta_net.cc	Implements CPU recurrence, fused gates, and state capture.

_{Review effort: Balanced}

[Copilot]

🟡 Changes recommended

Unresolved critical defects include out-of-bounds CUDA cache reads, silently incorrect scalar decay, and WebGPU shader-compilation and storage-binding failures.

24 open issues

• Validate initial-state rank before reading its batch dimension · New
• Reject nonmonotonic or out-of-range prefix offsets · New
• Validate prefix offsets before parallel processing · New
• Reject nonnegative gate lower bounds · New
• Honor the gate lower bound in scalar split and BF16 recomputation · New
• Preserve FP32 state range in the FP16 split path · New
• Preserve FP32 precision for the fp16 BF16-precision setting · New
• Retain the CUDA GEMV grid-y launch limit · New
• Widen FP16 params loads and include gate types in cache keys · New
• Support per-key decay in state replay or reject its captures · New
• Leave state-update width unknown when decay rank is unknown · New
• Pack KDA gate bindings to support adapter limits
• Convert half-precision gates and specialize pipeline metadata
• Guard cache loads for channels beyond head_size
• Handle saturated per-key decay prefixes without NaNs
• Propagate gate_lower_bound through split and BF16 paths
• Preserve legacy normalization epsilon outside KDA mode · New
• Size shared memory to the arrays actually carved · New
• Preserve a memory-safe fused FP8 prefill option · New
• Add WebGPU KDA, capture, and multi-pass test coverage · New

And 4 more that still need to be addressed.

_{Review effort: Max}

[Tianlei Wu (tianleiwu)]

Performance impact of the review-round changes, and a new option for the FP16 state guard

I measured the safety changes from the review round (3d71b25f95, c76e2f1a5a) against the original KDA kernels to see which ones cost real performance. Two did; for the others I found no cost that justifies an option.

Setup: H200 (sm_90), CUDA 13.0, Release build, KDA prefill (decay_per_key_dim, 128/128 heads) through the public API with device-resident IOBinding. Times are µs per call, median of 7 interleaved rounds of 40 calls.

1. Per-element decay finiteness check: restructured, no option

Checking every relative decay for a non-finite prefix cost 2.5–7.0% (fp16) and 1.9–3.9% (bf16). The prepare kernel now tests the last prefix sum of each key dim once per chunk (a non-finite decay makes every later prefix non-finite) and only a chunk that fails the test takes the checked path. The common path is the original single expf, within 1.3% of a kernel with no check at all, so this needs no option. Saturating-gate behavior is unchanged (KdaSplitFiniteSaturatingGates).

2. FP16 state range guard: real cost, so it is now optional

The scan keeps its recurrent state in FP32 and rounds it to fp16 for the tensor-core products. The guard computes the exact FP32 product for any tile that holds an element beyond 65504. Turning it off saves up to about 10% of the time for a single long sequence, and the saving shrinks as the batch grows:

fp16 shape	guard 1 (default)	guard 0	0 vs 1
B1 T256 H32	180.7	170.8	−5.5%
B1 T1024 H32	628.9	570.3	−9.3%
B1 T4096 H32	2295.4	2062.8	−10.1%
B1 T8192 H32	4496.1	4054.4	−9.8%
B4 T1024 H32	2250.5	2162.5	−3.9%
B8 T512 H32	4110.7	4096.2	−0.4%
B1 T4096 Hq16/Hv48	2734.4	2549.2	−6.8%
B1 T2048 Hq4/Hv128	3329.0	3131.4	−5.9%
B1 T4096 H32, fp32 gates	2271.2	2050.4	−9.7%

bf16 is unaffected (all nine bf16 shapes are within ±0.2%); it keeps ep.cuda.gated_delta_net_bf16_precision.

With an input whose state exceeds the fp16 range, the guarded kernel keeps the output and final state finite (max |out| 1.7e4), while the unguarded kernel returns 72% non-finite values in both. The two agree to 6e-5 where both are finite, and the unguarded kernel takes 13–14% less time on such input. Skipping the guard is therefore only safe when the state is known to stay in range, so the default stays guarded.

The option follows the other CUDA session switches (QMoE, MatMulNBits):

• session option ep.cuda.gated_delta_net_fp16_state_guard, environment variable ORT_GDN_FP16_STATE_GUARD
• "1" (default) guards, "0" skips the guard; any other value fails session initialization
• the session option overrides the environment variable and is resolved once at kernel creation

With the guard off, the kernels are as fast as (0.4–1.2% faster than) a variant with neither the guard nor the decay checks, which is the speed from before the review round. That comparison is an A/B build of the same kernels with the safety code removed, not a build of the original commit.

Other review-round changes

By inspection, the constructor and cu_seqlens validation and the CPU conv state checks are O(batch) work per call, the l2_eps select is uniform across the block, and the has_c0 predicate in the paged attention kernel only skips idle threads. None of them needs an option.

Not measured

WebGPU (not built here; the KDA shader changes are correctness fixes), CPU performance, and FP8 QMoE (not compiled in my build; it already has ORT_ENABLE_FP8_FUSED and ORT_QMOE_FP8_FUSED_MAX_ROWS).

Tests and docs

• New: GatedDeltaNetTest.KdaCudaFp16StateGuardSessionOptionAndEnv (default, environment variable, precedence, invalid values, unguarded overflow) and GatedDeltaNetTest.KdaCudaFp16StateGuardOffMatchesReference (fp16 and bf16, uniform and ragged).
• The GatedDeltaNet, causal-conv, gated-RMSNorm and linear-attention filters on the final build: 267 passed, 22 skipped (WebGPU tests, this build has no WebGPU EP), 0 failed. lintrunner is clean.
• docs/contrib_ops/cuda/gdn.md documents the option and its measured cost, and its engine-selection text is corrected: it still said KDA falls back to the sequential engine.

Pushed as 257da30fd7.

[Baiju Meswani (baijumeswani)]

This PR adds Kimi Delta Attention support to GatedDeltaNet on CPU, CUDA, and WebGPU, including per-key decay and compact state capture. It also adds a fused MoE group router and optimizes CUDA convolution, paged attention, and block-FP8 QMoE.

The direction looks good. Separating KDA preparation from the state scan is a useful way to expose more parallel work. Keeping recurrent state in FP32, retaining fallback paths, and adding reference-based tests across providers are also good choices. The long-prefill convolution optimization showed a substantial improvement in my focused A100 check.

I cannot recommend approval yet. The five inline comments describe two KDA numerical errors, missing router attribute validation, incorrect routing for larger expert groups, and a slowdown for short batched convolutions. Please address these and add focused regression tests. In particular, test the general router path as well as the specialized 128-expert path, and use deterministic numerical cases that make incorrect outputs clearly visible.

For maintainability, I would also consider separating the independent optimizations and moving MoEGroupRouter into its own implementation files and test suite. For further performance work, profile preparing shared normalized Q/K once when many value heads share each Q/K head. An explicit replay decay-layout attribute would also be clearer than inferring the layout from buffer-offset spacing. These are design suggestions, separate from the five concrete issues below.

Validation: I compiled the changed CUDA kernels into a standalone harness and ran the focused correctness and convolution timing checks on an A100. I checked the invalid CPU iterator with the checked C++ standard library. These checks were not a full ORT build or provider-suite run; I did not independently run WebGPU or end-to-end accuracy tests. Review is against 4f60f7b.