Attention, linear attention, KV cache writes, rotary embedding and normalization. The same op type can be the ONNX standard op (opset 23) or the com.microsoft op from ORT exporters, so the domain decides the input layout.
On EMLX (see Transmuter.Ops.emlx?/1), the normalizations and the attention core run as fused kernels of EMLX.Fast. The kernels take f16 and bf16 values as they are and accumulate in f32. Attention with a softcap keeps the Nx operations, because the kernel has no softcap. The Nx operations compute f16 and bf16 values in f32 (see Transmuter.Ops.widen/2).
All attention ops share one attention core with the scale, the softcap, the masks and the grouping of key and value heads. When the causal mask is the only mask, it runs as the causal kernel, which builds no mask. Other masks become one additive mask for the kernel.
Without the kernel, the attention core builds the scores of the query rows, in f32 or in the wider float type of q and k. When the scores of all rows take more than 256 MiB, the core runs on chunks of query rows, one chunk at a time. The rows split into the fewest chunks whose scores fit in 256 MiB, with the same number of rows in each chunk. A chunk has at least one row, so when the scores of one row take more than 256 MiB, the scores of a chunk take more too. Each row has its own softmax over all keys, so the chunks give the values of all rows at once. For example, a 1540x1540 image in the vision encoder of Ministral 3 has 12100 query rows and scores of 9.4 GB, which become 35 chunks of 346 rows and 268 MB.
AttentionCore (domain transmuter) is the attention core that Transmuter.Fusion finds in a graph: softmax(where(keep, q · k * scale / divisor, fill) + mask) · v, with the MatMul broadcasting of ONNX and the softmax along its axis attribute. Its inputs are q, k, v and the optional mask, keep, divisor and the two shapes of a Reshape pair of the probabilities. Its k is the right operand of the first MatMul. It runs as the shared attention core, the kernel on EMLX or the chunks of query rows, when q, k and v are 4-D heads of the same batch, or 3-D with the heads in the batch, the softmax is along the last axis, the masks broadcast to the scores, the host can read the divisor and the Reshape pair gives back the shape of the probabilities. The other layouts keep the plain operations of the graph and build the scores of all rows at once: the MatMul broadcasting, a divisor tensor, a softmax along another axis and a Reshape pair do not split by query row.