Custom operators

custom_op_schemas(op_type='', init_doc=True)

Returns LightOpSchema records for supported com.microsoft operators.

operator_schema_lookup(op_type)

Returns standard ONNX schemas, Microsoft custom schemas, and this package’s experimental compatibility schemas. Pass it as GraphBuilder(..., schema_lookup=operator_schema_lookup).

experimental_op_schemas(op_type='', init_doc=True)

Returns experimental ai.onnx compatibility LightOpSchema records. These adapters are separate from standardized ONNX schemas and from the Microsoft-only custom_op_schemas() provider.

Microsoft MultiHeadAttention

com.microsoft::MultiHeadAttention v1 accepts FLOAT, FLOAT16 and BFLOAT16 query (B,S,N*H), key (B,L,N*H) and value (B,L,N*Hv), or Whisper-style cross-attention key (B,N,L,H) and value (B,N,L,Hv). num_heads is required. The output shape is (B,S,N*Hv). Shape inference preserves the value’s symbolic total width for rank-3 inputs and the symbolic product N*Hv for rank-4 inputs. Optional inputs are projection bias (2*N*H+N*Hv), int32 key-padding lengths (B) or binary masks (B,T)/(B,S,T), additive attention bias (1 or B,1 or N,S,T), and paired past key (B,N,P,H) and value (B,N,P,Hv). For rank-4 cross attention, only the query segment of projection bias is applied; the key and value bias segments are ignored, as in the upstream CPU implementation. The optional present outputs contain P+L positions. scale defaults to 1/sqrt(H), mask_filter_value to -10000, and unidirectional to zero. Packed rank-5 inputs, shared-buffer inputs 8/9 and qk output are explicitly unsupported. The implementation is independent of ai.onnx::Attention and does not reinterpret its different mask and projection semantics. The contract is based on ONNX Runtime’s upstream schema and CPU implementation; the CPU kernel here is an original implementation, not a copy of upstream code.

The focused C++ tests cover scalar self/cross attention, projection bias, key padding, cached present outputs, all three floating types, symbolic shapes and the explicitly rejected slots. An end-to-end Whisper model run and direct numerical differential testing against ONNX Runtime remain untested; the backend fixtures use independently specified small expected outputs rather than importing ONNX Runtime.

Microsoft GroupQueryAttention

com.microsoft::GroupQueryAttention supports causal local attention for FLOAT, FLOAT16, and BFLOAT16. local_window_size defaults to -1 (full attention); a positive INT32 value counts attended tokens, including the current token. For example, 2048 attends the current position and up to 2047 preceding positions. Zero, values below -1, values above INT32_MAX, and local windows with causal=0 are rejected.

Windowing restricts only attention scores: present_key and present_value still contain the entire past cache followed by the current step. RoPE uses absolute positions as before. Full and sliding attention nodes can coexist in the same model; sliding_window_cache remains unsupported.

The optional twelfth input, head_sink, has shape (num_heads,) and the same element type as the query. Each query head’s sink is a virtual softmax logit: it contributes to the normalization denominator, but contributes no value to the output. Causal and local-window masks still apply only to real keys. Omitting the input preserves ordinary attention normalization.

Microsoft MoE

com.microsoft::MoE version 1 routes rank-2 packed tokens [T,H] or rank-3 decoder activations [B,S,H] using rank-2 [T,E] router logits. It applies a softmax over all experts, selects the top k (default 1), and optionally renormalizes the selected weights before summing expert outputs. Expert weights are [E,I,H] for FC1 and [E,H,I] for FC2. Optional expert-major biases occupy inputs 3 and 5. Interleaved SwiGLU (activation_type="swiglu", swiglu_fusion=1) uses FC1 [E,2*I,H] with alternating gate and linear rows, and supports swiglu_limit, activation_alpha, and activation_beta. The other supported activations are relu (default), gelu, silu, and identity.

Input and output types must match and be FLOAT, FLOAT16, or BFLOAT16. Low-precision inputs and weights are converted in bounded row tiles and accumulated in FLOAT before the output is narrowed. Separate FC3 weights and biases, sparse mixer routing, and noninterleaved SwiGLU are not supported by the CPU kernel. The output retains the input shape. See ONNX Runtime’s MoE CPU implementation for the upstream compatibility contract.

Microsoft QMoE

com.microsoft::QMoE version 1 accepts the block-wise CPU expert tensors emitted by mbext’s MoE builder. Activations, router logits, optional biases, and output share one of FLOAT, FLOAT16, or BFLOAT16. Scales may independently use any of these three floating-point types. Input/output and routing shapes, top-k softmax normalization, and activation parameters follow MoE above.

Inputs are, in order: input, router logits, FC1 weights, FC1 scales, optional FC1 bias, FC2 weights, FC2 scales, optional FC2 bias, optional FC3 weights, optional FC3 scales, optional FC3 bias, and three optional zero-point slots. FC3 weights and scales must appear together. Nonempty zero points, sparse mixer routing, and CUDA-prepacked weights are rejected.

expert_weight_bits supports 4 and 8. INT4 stores adjacent last-axis values in low-then-high nibbles with implicit zero point 8; INT8 stores signed two’s-complement values in UINT8 tensors for block quantization. block_size=0 (default) accepts only raw CPU output-major matrices with per-output-row scales [E,R]; INT8 then uses an implicit zero point of 128. Supported positive block sizes are powers of two of at least 16, including 16, 32, 64, and 128; scales then have shape [E,R,ceil(C/block_size)] for a stored logical matrix [E,R,C]. The final block and an odd INT4 column count need not be padded in the weight tensor: its shape is [E,R,ceil(C/2)] for INT4 or [E,R,C] for INT8.

For block quantization, the fused builder emits output-major FC1 [E,2*I,H] and FC2 [E,H,I] logical matrices with interleaved SwiGLU gate/up rows (activation_type="swiglu", swiglu_fusion=1). The split builder supplies FC3 and emits input-major FC1/FC3 [E,H,I] and FC2 [E,I,H] logical matrices. Its expert function is FC2(activation(FC1(input)) * FC3(input)), including optional biases. Separate FC3 and fused SwiGLU cannot be combined. Without FC3, ordinary relu, gelu, silu, and identity projections are also supported in input-major layout.

For block quantization, the kernel validates both orientations against all projection and scale shapes, preferring the builder’s orientation when dimensions are ambiguous. Tensor quantization always uses ONNX Runtime’s raw output-major, row-scaled format. See the ONNX Runtime CPU dequantizer.

This is a portable correctness implementation: expert matrices are dequantized into execution-allocator scratch and projections accumulate in FLOAT. It does not interpret CPU weights as CUDA’s MatMulNBits block packing or TensorRT’s shuffled prepacked format. Quantized weights are inference-only; no QMoE gradient or graph-fusion rule is registered.

The cited builder’s legacy qmoe_block_size=0 fallback selects TensorRT-LLM’s preprocessed quantizer output, not its raw CPU tensor output. That architecture-dependent CUDA format is not the raw tensor layout described above, even though it has the same tensor shape. Re-export those models with a positive CPU qmoe_block_size (for example, 32). Untagged CUDA-prepacked contents cannot be identified from tensor shapes alone; do not pass them as raw CPU tensors.

Microsoft GatherBlockQuantized

com.microsoft::GatherBlockQuantized version 1 gathers embeddings directly from packed UINT8 weights, without expanding the checkpoint to a dense floating-point table. bits supports 2, 4, and 8 (default 4). Codes occupy consecutive low-to-high bit fields within each byte. block_size defaults to 128 and must be a power of two of at least 16.

The supported layout has gather_axis=0 and quantize_axis on the last dimension (default 1); equivalent negative axes are accepted. Weights have rank at least two. If their stored shape is [V, ..., P], the logical last dimension is D = P * (8 / bits). Scales have the same rank and dimensions, except their last dimension is ceil(D / block_size). Their type, FLOAT or FLOAT16, determines the output type.

Optional UINT8 zero points use the same packing as the weights, with ceil(number_of_groups / (8 / bits)) bytes per last-axis row. Packing restarts for each row; unused fields in its final byte are ignored. When omitted, the zero point is 2 ** (bits - 1). Each selected code is dequantized as (code - zero_point) * scale. INT32 or INT64 indices may have any rank, including scalars. Negative indices wrap from the end; out-of-range indices produce zero embeddings, matching the ONNX Runtime CPU implementation. Older ORT releases, including 1.30, may instead reject out-of-range indices. The output shape is indices.shape + weights.shape[1:-1] + [D].

Register with register_kernels_global() before creating evaluators, or register the operator on the decoder session with register_kernel_for_session(decoder, "com.microsoft", "GatherBlockQuantized"). It can then execute in decoder.generate(..., embedding=embedding): the embedding model retains packed initializers and only selected rows are dequantized into the output. No heap scratch, dense-weight preparation, gradient rule, or graph fusion is registered. Logical INT4/UINT4 tensors and floating-point quantized formats are not supported.

See Packed embedding gather against ONNX Runtime for bounded ONNX Runtime parity checks and timings.

Microsoft MatMulNBits

com.microsoft::MatMulNBits version 1 multiplies floating-point activations by a block-quantized right-hand matrix without materializing the complete dequantized matrix. The CPU contract implements packed INT2, INT4, and INT8 formats with required K and N attributes, bits equal to 2, 4, or 8, a power-of-two block_size in [16, UINT32_MAX] (including 16, 32, 64, 128, and 256), accuracy_level 0 or 4, and weight_prepacked=0. Input B has shape [N, ceil(K / block_size), block_size * bits / 8] and scales have shape [N, ceil(K / block_size)] or the equivalent flat shape. Codes are packed low bits first within each byte; the last group is padded in storage and only its valid K elements contribute. Implicit zero points are 2, 8, and 128. Optional zero_points (input 3) may instead be packed UINT8 with shape [N, ceil(ceil(K / block_size) / (8 / bits))] (each column starts at a byte boundary), or floating-point values with the scales’ shape and dtype. Both zero-point layouts also accept the equivalent flat shape. This supports ternary INT2 groups of 128 with zero point 1 without changing their quantization geometry. Activations, scales, optional bias, and output have one matching FLOAT, FLOAT16, or BFLOAT16 type; DOUBLE is not supported.

Input B may instead be an EncodedValueProto using onnx-light’s ORT_MATMULNBITS_INT2, ORT_MATMULNBITS_INT4, or ORT_MATMULNBITS_INT8 layout. In that form, omit the scales input: the kernel validates the encoded geometry, dtype, layout, padding, and zero points, then borrows its packed B and scales without copying them per inference. It never converts the encoded weights to the VNNI preparation layout. Encoded layouts with shared parameter references, incompatible attributes, or any other unsupported representation fail with an exception. This encoded form is inference-only; automatic differentiation continues to require separate tensor B and scales inputs with block size 32 and implicit zero points. Explicit zero points are borrowed from the encoded payload; omit input 3 as well as input 2. Larger groups and explicit zero points use the bounded-panel decoder. The specialized INT4/INT8 preparation paths remain restricted to block size 32 with implicit zero points. g_idx is not supported.

Genuine three-bit weights use the separate onnx-light MATMULNBITS_INT3 format (struct name onnx_light.quantization.v1/matmulnbits_int3), not an ORT layout or a UINT8 tensor B. Set bits=3, block_size=32 and omit the scales node input. Its EncodedValueProto.logical_type is [K,N] with FLOAT, FLOAT16, or BFLOAT16 element type. The inline struct has, in order, an INT64[2] constant named parameters containing [3,32], a UINT8[N,ceil(K/32),12] field named B, a [N,ceil(K/32)] field named scales with the logical element type, and optionally a UINT8[N,ceil(K/32)] field named zero_points. Concatenate these fields (excluding the constant) in that order into raw_data; no external payload or shared parameters are permitted. Block order is output column first, then K block. In each 12-byte block, code q[p] occupies bits [3*p, 3*p+3) in little-endian bit order (least-significant bit first within each byte). Unused codes and bits after the final K element in each column must be zero. For example, codes [1,2,5] begin with bytes 51 01 (the low byte of 1 | (2 << 3) | (5 << 6) is 0x51; the third code’s high bit is in the second byte). The remaining ten bytes of this block are zero. An omitted zero-point field means code 4 for every block; explicit zero points must each be in [0,7]. Weights are reconstructed as (q - zero_point) * scale without expanding a full INT8 matrix. INT3 is not ternary: it represents eight distinct codes. The low-level C++ INT3 entry point also rejects overflowing activation, packed-weight, scale, and output sizes before dispatch or buffer access.

The activation may have any positive rank with final dimension K. Shape inference preserves every leading concrete or symbolic dimension and replaces the final dimension with N. INT4 uses bounded panels of 8 rows, 32 output columns, and 32 reduction elements, with scalar, AVX2, or AVX-512 dispatch. Scratch is 6,144 bytes per active callback (4,096 decoded weight bytes, 1,024 activation bytes, and 1,024 accumulator bytes), independent of matrix size. On AVX-512 VNNI/BW systems, FP32 INT4 with accuracy_level=4 keeps the weights packed at four bits in a 16-column VNNI layout, reorders the scales, and precomputes weight sums. This uses K*N/2 + 8*N*ceil(K/32) persistent bytes, dynamically quantizes each 32-value activation block to INT8, and accumulates in INT32 with vpdpbusd. Each active eight-row tile uses 8*K + 32*ceil(K/32) temporary bytes. Other combinations retain the bounded panel or allocation-free scalar path. Inputs and outputs are excluded from scratch accounting. The projection benchmark reports both worker-local and prepared storage.

Explicit zero points in the legacy layouts or node input, deprecated g_idx, provider-prepacked weights, other bit widths, block sizes, mixed floating-point types, and DOUBLE are rejected rather than silently using a different layout. The registered gradient differentiates the activation and optional bias while treating packed weights and scales as constants. The bias fusion applies only to an exclusively consumed MatMulNBits output followed by a compatible rank-one Add.

Muse-Glimmer projections

The synthetic projection suite uses M=1,8,128 for decode, short prompts, and prefill. Its K -> N families are Q/attention gate 6656 -> 4096, K/V 6656 -> 256, attention output 4096 -> 6656, MLP gate/up 6656 -> 19968, MLP down 19968 -> 6656, and vocabulary 6656 -> 202048. The vocabulary is the Transformers reference default, not a verified checkpoint value.

The mbext source plan does not supply an ONNX artifact. The generic exporter packs consecutive K values low-nibble first into UINT8 [N,ceil(K/G),G/2] and stores matching floating-point scales [N,ceil(K/G)]. Symmetric default/RTN export omits zero points (INT4 midpoint 8); G=32 is the configurable builder default. The activation dtype policy uses FP32 for CPU integer-weight export; CUDA normally uses FP16, or BF16 when requested. These are source-audited policies, not observations of a Muse model.

Separate unsupported contract: asymmetric and k_quant exports can provide explicit packed zero points [N,ceil(ceil(K/G)/2)]; exporters can also select a block size other than 32. Neither is enabled by this optimization. The actual artifact must be inspected for these attributes, zero-point inputs, dtypes, projection fusion, and whether the LM head is quantized before claiming model compatibility.

Run the reproducible parity/latency suite with python -m tools.benchmark_matmul_nbits_parity --help. It separates preparation from repeated invocations using constant initializers, compares against ONNX Runtime, and reports scratch/copy accounting separately from process memory. Pass --accuracy-level 4 to exercise the INT8-accumulation path. Large vocabulary cases are opt-in because model serialization and ONNX Runtime preparation can require several GB. For isolated single-thread FP32 kernel measurements, build with -DONNX_LIGHT_CPU_BUILD_BENCHMARKS=ON and run matmul_nbits_throughput M K N repeats [bits=4|8] [--generic]. This reports input preparation, first-use packing time, median steady-state seconds, source packed bytes, persistent prepared bytes, per-callback scratch bytes, and the selected implementation. For example, matmul_nbits_throughput 1 512 32000 21 8 --generic measures the generic INT8 panel path, not INT4. Both modes use one participant; the generic mode allocates no persistent prepared-weight buffers. The bit width remains optional with --generic: matmul_nbits_throughput 1 512 32000 21 --generic uses the default INT4 path.

The prepared INT8 path processes four activation rows per eight-column tile, reusing packed weights across rows and keeping dot products and accumulation inside the ISA-specific implementation. Activation blocks use the full signed INT8 range [-127,127]; AVX2 splits the unsigned activation’s high bit before pairwise multiplication to avoid 16-bit saturation. AVX-512 VNNI uses all 512 bits to process two weight quads at once.

Disjoint scratch of 5*K bytes per participant is allocated through the runtime execution allocator before dispatch; range callbacks do not allocate. For decode with M=1,2,3,4, activations are quantized once into shared, read-only scratch of 5*M*K/4 bytes and workers own disjoint groups of eight output columns. The parallel threshold uses the actual M*N output count, excluding padded lanes. For a zero reduction dimension, the low-level entry point returns the bias, or zero when no bias is supplied, without dispatch. The output-based parallel threshold, target block size, and participant cap apply to prepared execution as well as the generic path. Activation scales, scale products, and biased accumulation remain in double precision until each updated block accumulator is converted to FLOAT. Nonfinite activations or weight scales retain per-element floating-point dequantization semantics. Implementation labels remain available through used_kernel_paths(), but are excluded from used_kernel_names().

Run python -m tools.benchmark_matmul_nbits_parity --qwen-int8 for all six synthetic Qwen2 prefill projections at four threads. Both sessions use the same graph and input values; CPU timings feed native backend tensors without NumPy adaptation. Preparation and first-run times are separate from the alternating steady-state samples. These isolated nodes do not measure the exported end-to-end decoder. Use --qwen-int8-decode instead for all six projections at M=1,2,4 with one and four threads.

On an Intel Xeon Platinum 8480C with affinity restricted to four CPUs, GCC 13.2 Release and ONNX Runtime 1.32.0, the --qwen-int8 command with --repeat 15 --warmup 5 measured these median steady-state milliseconds (M=128, FP32, INT8, block size 32, accuracy level 4). The selected CPU implementation was int8_avx512_vnni; maximum absolute error against ORT across these six cases was below 1.9e-5.

Projection	K	N	CPU before (51e760e)	CPU tiled	ORT
LM head	512	32000	21.044	6.812	6.725
MLP gate	512	1376	0.900	0.337	0.401
MLP up	512	1376	0.900	0.325	0.421
MLP down	1376	512	0.951	0.328	0.484
Attention QKV	512	1024	0.704	0.251	0.337
Attention output	512	512	0.376	0.137	0.228

Both CPU revisions used native backend feeds, the same deterministic inputs, four threads, five warmups, and fifteen alternating samples. Preparation is excluded. The old seven-bit activation quantization also failed ORT parity on these cases (maximum LM-head error about 0.285); the timings above do not imply that its numerical behavior was equivalent. These measurements are not portable speedup thresholds, and the LM head remains slightly slower than ORT on this runner.

ONNX Runtime 1.30.0 CPU does not implement native BF16 MatMulNBits. The BF16 oracle therefore runs ORT’s FP32 operator on BF16-rounded activations and scales and rounds its output back to BF16. This validates the CPU BF16 path without duplicating the quantized operator in Python, but is not a claim of native ORT BF16 support.

Validation on this runner completed all 81 full-sized ORT comparisons: nine projection families at M=1,8,128, for all three activation types, including vocabulary prefill. The separate attention-gate case has the same dimensions and deterministic inputs as Q. Reproduce these checks through --full or the OLC_MUSE_INT4_FULL=1 test matrix. Maximum absolute errors across the completed comparisons were 4.816e-5 (FP32), 0.00390625 (FP16), and 0.015625 (BF16 oracle).

Representative kernel measurements

On an AMD EPYC 9V74 runner, GCC 13.3 Release, one runtime participant, FP32 K=6656,N=256, the standalone benchmark measured the following seconds (one warmup, median of five runs). The baseline is the foundation implementation at a16de32cd15d4c4ac3ba9d7644a92dc80b7c2e22; the optimized binary reported int4_panel_avx512. These are representative measurements, not portable performance thresholds.

M	Input preparation	Baseline steady state	Panel steady state	Panel scratch bytes	Packed-weight copy bytes/run
1	0.00101	0.00334	0.00154	6144	0
8	0.00114	0.0266	0.00213	6144	0
128	0.00285	0.425	0.0295	6144	0

Preparation above generates inputs and allocates output storage; the kernel has no packing/preparation phase. Python session construction and first-use costs are separate from these low-level timings. Packed weights occupy 851,968 bytes in this case; a full FP32 weight matrix would occupy 6,815,744 bytes. The same panel bound applies to the reference vocabulary projection, whose full FP32 weight matrix would occupy 5,379,325,952 bytes.

Experimental SimplifiedLayerNormalization

SimplifiedLayerNormalization is an experimental ONNX Runtime compatibility operator in the default domain ("", also spelled "ai.onnx"), since version 1. Use operator_schema_lookup() to load or incrementally construct its graph with GraphBuilder.make_node. Register shape and memory support with register_operator_support(); all global and session-local kernel registration helpers also register this support.

X and Scale independently support FLOAT, FLOAT16, DOUBLE, and BFLOAT16 (all sixteen pairs). Mandatory Y has the shape of X and the element type of Scale. Scale broadcasts right-aligned to the entire shape of X, not only the normalized suffix, and cannot expand the shape of X. axis=-1 selects the first normalized dimension and epsilon=1e-5 is added to the mean square. IEEE epsilon values, including negative and non-finite values, are permitted.

The second output, inv_std_var, is optional and may be omitted or named with an empty string. stash_type=1 saves it as FLOAT; stash_type=11 saves it as DOUBLE. Other stash types are rejected. Its runtime-compatible shape is X[:axis] + [1] * (rank(X) - axis) after normalizing negative axis: every normalized dimension is one, unlike upstream schema inference which only replaces the axis dimension. X must have positive rank and a nonempty normalized suffix; empty outer rows are supported.

stash_type selects arithmetic precision as well as the statistics type: 1 uses FP32 and 11 uses FP64, independently of input types, broadcasting, and whether statistics are requested. Inputs and scale are converted to this precision for the reduction, normalization, and scaling; the result is then converted to the type of Scale. In particular, FP64 inputs with the default FP32 stash can overflow when converted or squared. There is no intermediate low-precision rounding before multiplication by Scale.

The FP32 suffix path reuses the optimized RMS mean-square and affine engine. FP16 with FP32 stash uses shared F16C primitives with a single final rounding; FP64 uses AVX primitives for either stash type. Contiguous scale suffixes use one broadcast index per row or block, including outer-broadcast scales. Unsupported SIMD/type/alignment combinations retain portable fallbacks. Optional statistics are saved without repeating the reduction. Benchmark cpu_kernel_paths diagnostics distinguish the RMS, FP16/F16C, FP64/AVX, scalar, and generic paths, together with their scale layout.

CPU tensor scratch memory is zero, excluding inputs and outputs (broadcast metadata is proportional to rank). This is inference compatibility support: no gradient rules or fusion patterns are registered for this operator.

Microsoft SkipLayerNormalization

com.microsoft::SkipLayerNormalization version 1 adds input, skip, and optional bias before centering the residual sum over its last axis, scaling by the inverse standard deviation and gamma, and adding optional beta. All inputs have one matching FLOAT, FLOAT16, or BFLOAT16 type. input has shape [S, H] or [B, S, H]; for rank-three input, skip may also have shape [S, H] or [1, S, H]. gamma, optional beta, and optional bias are rank-one tensors of length H.

Output slot 0 has the input shape and type. Optional slots 1 and 2 are FLOAT mean and inverse-standard-deviation tensors with final dimension one. Optional slot 3 is the typed residual sum. epsilon defaults to 1e-12 and must be finite and nonnegative. The allocation-free CPU kernel uses FP32 arithmetic for narrow inputs and runtime-owned row scheduling. This follows the ONNX Runtime CPU contract. Shape inference preserves symbolic dimensions and scratch-memory accounting is zero, excluding inputs and outputs.

Microsoft SkipSimplifiedLayerNormalization

com.microsoft::SkipSimplifiedLayerNormalization version 1 adds input, skip, and optional bias, then normalizes the residual sum by sqrt(mean(sum * sum, axis=-1, keepdims=True) + epsilon) and multiplies by gamma. All inputs have the same FLOAT, FLOAT16, or BFLOAT16 type. input has rank two or three, with shape [S, H] or [B, S, H]. skip matches that shape; for rank-three input only, [S, H] and [1, S, H] also broadcast over the batch dimension. Sequence and hidden dimensions do not broadcast. gamma and optional bias are rank-one tensors of length H. 0 < H <= INT_MAX; empty outer dimensions are valid. epsilon defaults to 1e-12 and must be finite and nonnegative.

Mandatory output slot 0 has the input shape and type. Optional output slot 3, input_skip_bias_sum, contains the residual sum with the same shape and type. Request it with the ONNX output list ["Y", "", "", "input_skip_bias_sum"]. Optional slot 1, mean, contains FLOAT zeros; optional slot 2, inv_std_var, contains the FLOAT inverse RMS. Both statistics have the input shape with the final dimension replaced by one. They are independently optional; for example, ["Y", "", "inv_std_var"] requests only the inverse RMS. The bias may be omitted or represented by an empty input name. Trailing optional output slots may likewise be omitted or left empty.

Arithmetic matches the ONNX Runtime CPU contract: the inputs are widened to FP32 before residual addition, reduction, normalization, and scaling. FLOAT16 and BFLOAT16 values are narrowed only when storing the final normalized and residual outputs, never before computing the mean square. Statistics use this unrounded FP32 residual. The mean output is zero because RMS normalization does not subtract a mean. This follows the pinned ONNX Runtime CPU implementation; older ONNX Runtime versions and GPU implementations may not support these statistics or BFLOAT16.

Use custom_op_schemas() or operator_schema_lookup() as the native GraphBuilder schema lookup callback and register_operator_support() to register shape and memory support. Global and session-local kernel registration helpers also register this metadata. Shape inference preserves symbolic dimensions and records equality and batch-broadcast constraints. CPU scratch memory is zero, excluding inputs and outputs: the FLOAT path uses the normalized output buffer for the residual intermediate, while low-precision paths recompute the widened residual rather than allocating a temporary buffer. Saved statistics do not imply automatic differentiation support: no gradient rules or fusion patterns are registered.

Support inventory

class OperatorSupport

Immutable NamedTuple describing shape inference, peak memory, fusion patterns, and gradient availability for one custom or experimental operator.

operator_support() → tuple[OperatorSupport, ...]

Returns the custom and experimental operator support inventory without registering or executing an implementation.

register_operator_support() → None

Registers custom/experimental shape and peak-memory support and custom fusion patterns.

register_custom_gradients(registry=None)

Adds the CDist, BiasGelu, GroupQueryAttention, and MatMulNBits backward rules to an independent GradRegistry and returns it.