Custom operators#
The custom operator support inventory and its public implementations are available through these APIs.
-
std::vector<OperatorSupportRegistration> onnx_light_cpu::CollectOperatorSupport()#
Returns the custom operator support available from onnx-light-cpu.
-
void onnx_light_cpu::RegisterMicrosoftShapeAndMemoryFunctions()#
Registers the custom shape-inference and peak-memory functions.
-
void onnx_light_cpu::RegisterCustomOperatorGradients(ONNX_LIGHT_NAMESPACE::core::gradient::GradRegistry ®istry)#
Registers standard-ONNX gradient graphs for
com.microsoftoperators.BiasGelu gradient:
┌─────┐ A, B ─────→│ Add │────→ z └─────┘ │ │ ┌──────────┐ ├───→│ CDF term │────→ cdf ──────────────────────────┐ │ └──────────┘ │ │ │ │ ┌──────────────────┐ │ ├───→│ Gaussian density │────→ density │ │ └──────────────────┘ │ │ │ ↓ ↓ └────────────────────────────→┌─────┐ ┌─────┐ │ Mul │────────────→│ Add │───→ derivative └─────┘ └─────┘ │ ↓ ┌─────┐ dC ────────────────────────────────────────────────────────────→│ Mul │───→ dz ───→ dA └─────┘ │ ↓ ┌────────────────────────┐ │ ReduceSum leading axes │───→ dB └────────────────────────┘The CDF term is
0.5 * (1 + erf(z / sqrt(2)))and the Gaussian density isexp(-0.5 * z^2) / sqrt(2 * pi). Their combination givesderivative = CDF term + z * density.CDist gradient:
┌───────┐ A, B ───→│ CDist │────→ C └───────┘ │ ↓ ┌───────────────────┐ dC ─────────────────→│ Metric derivative │────→ pair gradient └───────────────────┘ │ ↓ ┌──────────────┐ │ Unsqueeze(2) │───────────┐ └──────────────┘ │ │ ┌──────────────┐ │ A ────→│ Unsqueeze(1) │───┐ │ └──────────────┘ │ │ │ │ ┌──────────────┐ │ ┌─────┐ │ B ────→│ Unsqueeze(0) │───┴──→│ Sub │────→ difference ───────────────┐ │ └──────────────┘ └─────┘ │ │ ↓ ↓ ┌─────┐ │ Mul │ └─────┘ │ ┌─────────────────────────┤ │ │ ↓ ↓ ┌──────────────┐ ┌──────────────┐ │ ReduceSum(1) │───→ dA │ ReduceSum(0) │ └──────────────┘ └──────────────┘ │ ↓ ┌─────┐ │ Neg │───→ dB └─────┘The metric derivative is
dC / Cfor Euclidean distance (zero whereC == 0), and2 * dCfor squared Euclidean distance.GroupQueryAttention gradient:
Q, K, V ────→┌────────────┐────→┌────────────────────┐────→ probabilities │ Cast FLOAT │ │ Attention (mode 3) │ └────────────┘ └────────────────────┘ dY ─────────→┌────────────┐────→ float dY │ Cast FLOAT │ └────────────┘ probabilities, float dY, expanded V │ ↓ ┌──────────────────┐ │ Softmax backward │ └──────────────────┘ │ ↓ ┌───────────────┐ │ scale dScores │ └───────────────┘ │ │ ↓ ↓ ┌───────────┐ ┌───────────┐ │ MatMul K │ │ MatMul Q │ └───────────┘ └───────────┘ │ │ ↓ ↓ ┌────────────┐ ┌───────────────┐ │ reshape dQ │ │ reduce groups │ └────────────┘ └───────────────┘ │ │ │ ↓ ↓ ↓ ┌─────────────┐ ┌─────────────┐ ┌─────────────┐ │ CastLike(Q) │ │ CastLike(K) │ │ CastLike(V) │ └─────────────┘ └─────────────┘ └─────────────┘ │ │ │ ↓ ↓ ↓ dQ dK dVThe score scale is replayed from the forward attribute. When omitted, the graph computes
1 / sqrt(Shape(Q)[2] / num_heads)dynamically. Grouped K/V head gradients are summed back to their original head count.
-
class BiasGeluFusionPattern : public ONNX_LIGHT_NAMESPACE::core::builder::PatternOptimization#
Fuses an exact Gelu applied to an Add with a rank-one bias.
Before: ┌─────┐ A, bias ─────→│ Add │ └─────┘ │ ↓ z │ ↓ ┌──────┐ │ Gelu │────→ Y └──────┘ After: ┌──────────┐ A, bias ─────→│ BiasGelu │────→ Y └──────────┘The bias must broadcast over the last dimension of a floating-point input, and the Add output must be consumed exclusively by the exact Gelu.
-
class CDistFusionPattern : public ONNX_LIGHT_NAMESPACE::core::builder::PatternOptimization#
Fuses the squared-Euclidean distance expansion into CDist.
Before: ┌──────────────┐ A ────→│ Unsqueeze(1) │───┐ └──────────────┘ │ │ ┌──────────────┐ │ ┌─────┐ B ────→│ Unsqueeze(0) │───┴──→│ Sub │────→ difference └──────────────┘ └─────┘ │ ↓ ┌──────────┐ │ Mul self │ └──────────┘ │ ↓ squared difference │ ↓ ┌───────────────┐ │ ReduceSum(-1) │────→ Y └───────────────┘ After: ┌──────────────────────────┐ A, B ─────→│ CDist metric=sqeuclidean │────→ Y └──────────────────────────┘Both inputs must have rank two. The difference must be squared by an exclusively consumed self-Mul and reduced over the last axis without keeping the reduced dimension.
-
class GroupQueryAttentionFusionPattern : public ONNX_LIGHT_NAMESPACE::core::builder::PatternOptimization#
Rewrites the rank-3 grouped-query subset of
ai.onnx::Attentiontocom.microsoft::GroupQueryAttention.Before: ┌─────────────────────────┐ query, key, ─────→│ Attention │────→ output value │ q_heads != kv_heads │ └─────────────────────────┘ After: key ─────────────────→┌───────┐────→┌────────┐────→ sequence length │ Shape │ │ Gather │ └───────┘ └────────┘ │ └──→┌─────┐────→┌──────┐──┐ │ Sub │ │ Cast │ │ └─────┘ └──────┘ │ ↓ query ───────────────→┌───────┐────→┌────────┐────→┌───────────┐────→┌────────┐ │ Shape │ │ Gather │ │ Unsqueeze │ │ Expand │ └───────┘ └────────┘ └───────────┘ └────────┘ │ ↓ ┌─────────────────────────────────────────────┐ seqlens_k query, key, ─────→│ com.microsoft::GroupQueryAttention │←────┘ value │ num_heads, kv_num_heads, causal, scale │────→ output └─────────────────────────────────────────────┘The generated shape subgraph derives the cache-free sequence metadata required by GroupQueryAttention without changing the Attention result.
-
class LinearAttentionFusionPattern : public ONNX_LIGHT_NAMESPACE::core::builder::PatternOptimization#
Rewrites the FLOAT subset of
ai.onnx::LinearAttention-27to the compatiblecom.microsoft::LinearAttention-1contract.Before: query, key, value --->+--------------------------+---> output past, decay, beta --->| ai.onnx::LinearAttention |---> present_state +--------------------------+ After: query, key, value --->+----------------------------------+---> output past, decay, beta --->| com.microsoft::LinearAttention-1 |---> present_state +----------------------------------+
Inputs, optional-input slots, outputs, and all semantic attributes are preserved. The Microsoft-domain import is added to the graph.