[CUDA] Add GatedLatentPool and SparseRowSelect contrib operators

[Tianlei Wu (tianleiwu)]

Description

Adds two CUDA contrib operators that together form the compression and selection front end of a sparse-attention stack: GatedLatentPool pools ratio consecutive tokens into one latent cache row, and SparseRowSelect scores those rows against the current query and emits the topk row ids attention should gather. DeepSeek-V4 uses this pair (its KV compressor and its Lightning Indexer), but the schemas are written against the general pattern — gated pooling into a strided latent cache, then a lightweight multi-head scorer over that cache.

As ONNX primitives each of these is a long chain of gather/softmax/reduce/rotate nodes whose intermediates are far larger than their inputs; fusing them removes both the launches and the materialization.

Summary of Changes

GatedLatentPool

File	Change
onnxruntime/contrib_ops/cuda/math/gated_latent_pool.{h,cc}	OpKernel — attributes, validation, output shapes
onnxruntime/contrib_ops/cuda/math/gated_latent_pool_impl.{h,cu}	Fused pooling, RMS norm, rotary and optional quantization-simulation kernels

Gated pooling of ratio consecutive tokens into one latent row:

w      = Softmax(Where(valid, full_score[window] + ape, -1e30), over the window)
pooled = ReduceSum(full_kv[window] * w, over the window)
normed = norm_weight * pooled / Sqrt(ReduceMean(pooled * pooled, -1) + epsilon)

Design points worth a reviewer's attention:

• The operator carries a rolling state of the window_multiplier * ratio projections that preceded this step, so a row whose window straddles a step boundary can still be pooled without the caller re-feeding history.
• score is optional: when omitted, kv is twice as wide and carries the value in the low half of each row and the gate in the high half, which is what a single fused projection GEMM produces.
• rows is sized (seq_len - 1) / ratio + 2 — one row per slot the step can touch plus a spare — so the output shape never depends on past_lens, which keeps the op CUDA-graph-capturable. first_slot / last_slot / row_count say which of those rows are live.
• simulate_fp8 and simulate_rotated_fp4 apply a simulated FP8-E4M3 / Hadamard-rotated FP4-E2M1 round trip, matching what a quantized attention cache will store.

SparseRowSelect

File	Change
onnxruntime/contrib_ops/cuda/math/sparse_row_select.{h,cc}	OpKernel — attributes, validation, output shapes
onnxruntime/contrib_ops/cuda/math/sparse_row_select_impl.{h,cu}	Rotary + optional FP4 simulation, cache update, scoring and top-k kernels

score[b, s, c] = sum_h Relu(dot(query[b, s, h], present_cache[b, c])) * weights[b, s, h] * scale

• Rows produced by GatedLatentPool this step are written at first_slot + j; slots below keep past_cache. Slots above last_slot are deliberately left unspecified — no step has ever written them and no query can reach them. On a long-context export they are almost the whole cache, so producing them would dominate the operator.
• A query at absolute position past_lens[b] + s may only see the first (past_lens[b] + s + 1) / ratio rows.
• selection is emitted as row_id_offset + c so it can be concatenated straight into a paged cache's index list, padded with -1 when fewer than topk rows are visible.
• Attention does not depend on gather order, so the selection is emitted in ascending row order rather than by descending score — this avoids a sort and makes the output deterministic.

Shared

File	Change
onnxruntime/contrib_ops/cuda/math/quant_sim_common.cuh	Shared quantization-simulation helpers (QuantSimConv, block scale, E4M3 rounding, FP4 rotation) used by both operators
onnxruntime/core/graph/contrib_ops/contrib_defs.cc	Schemas, shape inference and the formulas above
onnxruntime/core/graph/contrib_ops/ms_opset.h	Schema registration
onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc	CUDA kernel registration for float, float16 and bfloat16
docs/ContribOperators.md, docs/OperatorKernels.md	Regenerated entries

Testing

• onnxruntime/test/python/transformers/test_gated_latent_pool.py
• onnxruntime/test/python/transformers/test_sparse_row_select.py

Both compare against PyTorch references built from the primitive formulas, across float32/float16/bfloat16, with and without the optional score input, with and without each quantization-simulation mode, and across window multipliers, ratios, and prefill/decode step shapes.

python -m pytest onnxruntime/test/python/transformers/test_gated_latent_pool.py \
                 onnxruntime/test/python/transformers/test_sparse_row_select.py
# 11 passed, 138 subtests passed

Additive only — no existing operator, schema or kernel is modified.

Checklist

• Tests added
• No breaking changes
• Documentation updated

[Copilot]

Pull request overview

Adds CUDA contrib operators for sparse-attention compression and row selection.

Changes:

• Implements fused gated pooling, cache updates, scoring, and top-k selection.
• Adds shared FP8/FP4 quantization simulation.
• Registers, documents, and tests both operators.

Reviewed changes

Copilot reviewed 16 out of 16 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
onnxruntime/test/python/transformers/test_sparse_row_select.py	Adds numeric parity tests.
onnxruntime/test/python/transformers/test_gated_latent_pool.py	Adds pooling parity tests.
onnxruntime/core/graph/contrib_ops/ms_opset.h	Registers schemas.
onnxruntime/core/graph/contrib_ops/contrib_defs.cc	Defines schemas and shape inference.
onnxruntime/contrib_ops/cuda/math/sparse_row_select.h	Declares the selection kernel.
onnxruntime/contrib_ops/cuda/math/sparse_row_select.cc	Validates and dispatches selection.
onnxruntime/contrib_ops/cuda/math/sparse_row_select_impl.h	Defines selection parameters.
onnxruntime/contrib_ops/cuda/math/sparse_row_select_impl.cu	Implements CUDA scoring and selection.
onnxruntime/contrib_ops/cuda/math/quant_sim_common.cuh	Adds shared quantization helpers.
onnxruntime/contrib_ops/cuda/math/gated_latent_pool.h	Declares the pooling kernel.
onnxruntime/contrib_ops/cuda/math/gated_latent_pool.cc	Validates and dispatches pooling.
onnxruntime/contrib_ops/cuda/math/gated_latent_pool_impl.h	Defines pooling parameters.
onnxruntime/contrib_ops/cuda/math/gated_latent_pool_impl.cu	Implements fused CUDA pooling.
onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc	Registers CUDA kernels.
docs/OperatorKernels.md	Lists kernel support.
docs/ContribOperators.md	Documents both operators.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.