include_preview_kernels.h

namespace onnx_light

Alias that makes onnx-light headers compatible with code that references ONNX_LIGHT_NAMESPACE (the macro used in the standard onnx package).

Set to ONNX_LIGHT_NAMESPACE so both names resolve to the same namespace.

Symbol-visibility attribute for the public onnx-light C++ API.

Defined as empty because onnx-light does not require explicit __declspec(dllexport) or __attribute__((visibility("default"))) annotations — visibility is controlled at the shared-library level. The macro is provided so that vendored ONNX headers that decorate their declarations with ONNX_API compile without modification.

Namespace alias so that ONNX C++ code (and consumers such as onnxruntime) that refers to the literal onnx namespace — rather than the ONNX_NAMESPACE macro — resolves to the onnx-light namespace. The standard onnx package lives in namespace onnx; onnx-light uses onnx_light (via ONNX_LIGHT_NAMESPACE), so this alias keeps onnx-light a true drop-in. It is only introduced when the onnx-light namespace differs from onnx.

namespace onnx_kernels

namespace kernel

class FlexAttention : public onnx_light::onnx_kernels::kernel::KernelBase

#include <include_preview_kernels.h>

Reference implementation of ai.onnx.preview::FlexAttention (v1).

Computes Y = Softmax((Q @ K^T) * scale, axis=-1) @ V over rank-4 (batched, multi-head) floating-point inputs (FLOAT or DOUBLE; FLOAT16 and BFLOAT16 are supported via internal promotion to FLOAT32). Supports Grouped Query Attention (GQA): when q_num_heads != kv_num_heads each K/V head is shared by a contiguous group of query heads, i.e. query head h uses K/V head floor(h / (q_num_heads / kv_num_heads)); q_num_heads must be a multiple of kv_num_heads.

The optional score_mod modifier subgraph is supported via the ScoreModFn callback overload: callers that have a way to evaluate the score_mod subgraph (e.g. a graph-executor) can pass a callable that receives the pre-softmax score tensor of shape (batch_size, q_num_heads, q_seq_len, kv_seq_len) and rewrites it in place; the kernel then runs Softmax on the rewritten scores before the final probs @ V matmul. Overloads without a callback retain the baseline behavior (score_mod treated as identity).

The optional prob_mod modifier subgraph is supported via the ProbModFn callback overload: callers that have a way to evaluate the prob_mod subgraph (e.g. a graph-executor) can pass a callable that receives the post-softmax probability tensor of shape (batch_size, q_num_heads, q_seq_len, kv_seq_len) and rewrites it in place; the kernel then uses the rewritten probabilities for the final probs @ V matmul. Overloads without a callback retain the baseline behavior (prob_mod treated as identity).

FLOAT and DOUBLE tensors are supported (and FLOAT16/BFLOAT16 via internal promotion to FLOAT32). Q, K and V must share the same element type.

Public Types

using ScoreModFn = std::function<void(Tensor&)>

Callback used to apply the score_mod modifier subgraph to the pre-softmax score tensor of shape (batch_size, q_num_heads, q_seq_len, kv_seq_len). The callback receives a mutable reference to a tensor sharing Q’s element type and must rewrite its contents in place while preserving data_type and shape; the kernel validates these invariants after the call.

using ProbModFn = std::function<void(Tensor&)>

Callback used to apply the prob_mod modifier subgraph to the post-softmax probability tensor of shape (batch_size, q_num_heads, q_seq_len, kv_seq_len). The callback receives a mutable reference to a tensor sharing Q’s element type and must rewrite its contents in place while preserving data_type and shape; the kernel validates these invariants after the call.

Public Functions

Tensor operator()(const Tensor &Q, const Tensor &K, const Tensor &V) const

Computes the attention output for the given Q, K, V tensors using the default scaling factor 1 / sqrt(head_size).

Tensor operator()(const Tensor &Q, const Tensor &K, const Tensor &V, float scale) const

Computes the attention output for the given Q, K, V tensors using an explicit scale value (matching the scale attribute of the operator).

Tensor operator()(const Tensor &Q, const Tensor &K, const Tensor &V, float scale, const ProbModFn &prob_mod) const

Computes the attention output and applies the prob_mod callback to the post-softmax probability tensor before computing Y = probs @ V. When prob_mod is an empty std::function this is equivalent to the overload without a prob_mod callback.

Tensor operator()(const Tensor &Q, const Tensor &K, const Tensor &V, float scale, const ScoreModFn &score_mod, const ProbModFn &prob_mod) const

Computes the attention output and applies the score_mod callback to the pre-softmax score tensor before the softmax, and the prob_mod callback to the post-softmax probability tensor before computing Y = probs @ V. When either callback is an empty std::function it is treated as identity.

void operator()(const Tensor &Q, const Tensor &K, const Tensor &V, float scale, Tensor &output) const

In-place overload writing into a caller-allocated output tensor. output must already share Q’s element type and have shape (batch_size, q_num_heads, q_seq_len, v_head_size) with its data buffer sized to match.

void operator()(const Tensor &Q, const Tensor &K, const Tensor &V, float scale, const ProbModFn &prob_mod, Tensor &output) const

In-place overload with prob_mod support. See the returning overload taking a ProbModFn for the semantics of prob_mod; output has the same preconditions as the in-place overload without a callback.

void operator()(const Tensor &Q, const Tensor &K, const Tensor &V, float scale, const ScoreModFn &score_mod, const ProbModFn &prob_mod, Tensor &output) const

In-place overload with score_mod and prob_mod support. See the returning overloads for the semantics of the callbacks; output has the same preconditions as the in-place overload without a callback.

inline explicit KernelBase(const KernelContext &ctx)

Public Static Functions

static inline constexpr bool CanRunInPlace() noexcept

FlexAttention computes a fresh output buffer from independent reads of Q, K, V and never aliases an input buffer.