com.microsoft.GroupQueryAttention

com.microsoft · ONNX Runtime contrib operator · contrib since_version 1

Description

Grouped-query attention for explicit BSH Q/K/V and BNSH caches. Direct Q/K/V supports bidirectional attention or causal local windows and may store its generated float cache independently as float16 or float32; existing unquantized cache inputs match the Q/K/V dtype. Causal cache paths support rotary embeddings, sliding windows, bias, head sinks, softcap, smooth softmax, and paired Q/K RMS normalization. Int8/int4 caches require float32 Q/K/V and output; int4 is prompt-only. Packed QKV, position IDs, interleaved rotary, bfloat16/float8, and diagnostic QK output are not implemented.

See the ONNX Runtime GroupQueryAttention contrib-operator spec for the reference semantics.

Inputs

Outputs

Attributes

Attributes and default values (overridable per request):

Type constraints

Implementation variants

One implementation is selected per call from the device capabilities, the request shapes and the dtypes; these notes say what each one covers.

• qkv_present_materialized_sgmat_f32 — Materialized float32 subgroup-matrix prefill for bidirectional QKV without past state. seqlens_k is metadata-only; the score pass emits per-row softmax statistics, the apply pass normalizes and applies the scores, and the present copy preserves the supplied key and value tensors.
• past_kv_materialized_sgmat_f32 — Materializes the causal score matrix with float32 subgroup-matrix tiles over the shared float cache, then applies the softmax-normalized weights with subgroup-matrix tiles. Both passes skip key tiles beyond each query tile's causal bound.
• past_kv_bias_materialized_sgmat_f32 — Materializes causal scores with float32 subgroup-matrix tiles, broadcasts and adds the attention bias across batch and head dimensions, folds the biased scores into row statistics, and applies the normalized weights with subgroup-matrix tiles.
• past_kv_materialized_sgmat_f16 — Materializes the causal score matrix with f16 operand tiles feeding f32-accumulating subgroup matrices over the shared cache and applies the softmax-normalized weights the same way; scores and row statistics stay f32. The score pass skips key tiles past each query tile's causal bound and the apply pass stops its reduction there.
• past_kv_rotary_materialized_sgmat_f32 — Rotates each query block at its absolute positions, then materializes causal scores against rotary-transformed cached keys and applies the normalized weights. The score and apply passes skip tiles beyond each query tile's causal bound.
• past_kv_rotary_materialized_sgmat_f16 — Materializes the causal score matrix with f16 operand tiles feeding f32-accumulating subgroup matrices over the shared cache and applies the softmax-normalized weights the same way; scores and row statistics stay f32. The score pass skips key tiles past each query tile's causal bound and the apply pass stops its reduction there.
• new_kv_past_materialized_sgmat_f32 — For chunked prefill, merges the past cache with new key/value rows, materializes causal scores over the merged cache, and applies the normalized weights using the same right-aligned causal bound.
• window_shift_materialized_sgmat_f32 — Materializes causal chunked-prefill attention with a windowed cache. The shift pass compacts surviving rows and appends the chunk; score and apply passes enforce both the sliding-window floor and causal bound.
• share_append_materialized_sgmat_f32 — Materializes causal chunked-prefill attention while updating a shared-capacity cache in place. Score and apply passes bound every causal tile by the live length from seqlens_k, not the buffer capacity, so right-padded batches remain left-aligned.
• new_kv_share_append_split — Retains the existing cache and appends new key/value rows in separate passes before portable attention. It avoids rebuilding unchanged cache positions when the past and present allocations share capacity.
• qkv_present_tiled_nosg — Portable tiled prefill route that computes attention online and writes the present cache separately. It is used when the flash shape is valid but no suitable subgroup route is admissible.
• quant_int8_decode_splitk — Appends int8-quantized key/value rows, partitions cached decode across the key axis, and merges partial online-softmax results. Used when one workgroup per head exposes too little independent work.
• qkv_present_flash_splitk — Partitions direct Q/K/V attention across the key axis, merges partial online-softmax results, and writes the present cache separately. It serves short-query shapes needing more key-axis parallelism.
• qkv_present_flash_cluster — Computes clustered online-softmax prefill directly from Q/K/V and writes the present cache separately. The family provides subgroup and portable reductions for the same tiled algorithm.
• quant_int8_decode_splitk_nosg — Appends int8-quantized key/value rows, partitions cached decode across the key axis, and merges partial online-softmax results. Used when one workgroup per head exposes too little independent work.
• qkv_present_flash_splitk_nosg — Partitions direct Q/K/V attention across the key axis, merges partial online-softmax results, and writes the present cache separately. It serves short-query shapes needing more key-axis parallelism.
• qkv_present_flash_cluster_nosg — Computes clustered online-softmax prefill directly from Q/K/V and writes the present cache separately. The family provides subgroup and portable reductions for the same tiled algorithm.
• past_kv_decode_splitk — Updates the float cache in copy, merge, or window-shift mode, then partitions cached decode across the key axis and combines partial online-softmax results. The family includes subgroup and portable forms.
• new_kv_past_decode_splitk — Updates the float cache in copy, merge, or window-shift mode, then partitions cached decode across the key axis and combines partial online-softmax results. The family includes subgroup and portable forms.
• window_shift_decode_splitk — Updates the float cache in copy, merge, or window-shift mode, then partitions cached decode across the key axis and combines partial online-softmax results. The family includes subgroup and portable forms.
• past_kv_decode_splitk_nosg — Updates the float cache in copy, merge, or window-shift mode, then partitions cached decode across the key axis and combines partial online-softmax results. The family includes subgroup and portable forms.
• new_kv_past_decode_splitk_nosg — Updates the float cache in copy, merge, or window-shift mode, then partitions cached decode across the key axis and combines partial online-softmax results. The family includes subgroup and portable forms.
• window_shift_decode_splitk_nosg — Updates the float cache in copy, merge, or window-shift mode, then partitions cached decode across the key axis and combines partial online-softmax results. The family includes subgroup and portable forms.
• past_kv_flash_prefill — Updates the float cache in the selected mode and applies clustered online-softmax causal prefill with the configured cache, mask, rotary, softcap, bias, or head-sink features.
• new_kv_past_flash_prefill — Updates the float cache in the selected mode and applies clustered online-softmax causal prefill with the configured cache, mask, rotary, softcap, bias, or head-sink features.
• past_kv_rotary_flash_prefill — Updates the float cache in the selected mode and applies clustered online-softmax causal prefill with the configured cache, mask, rotary, softcap, bias, or head-sink features.
• past_kv_softcap_flash_prefill — Updates the float cache in the selected mode and applies clustered online-softmax causal prefill with the configured cache, mask, rotary, softcap, bias, or head-sink features.
• past_kv_headsink_flash_prefill — Updates the float cache in the selected mode and applies clustered online-softmax causal prefill with the configured cache, mask, rotary, softcap, bias, or head-sink features.
• past_kv_bias_headsink_flash_prefill — Updates the float cache in the selected mode and applies clustered online-softmax causal prefill with the configured cache, mask, rotary, softcap, bias, or head-sink features.
• window_shift_flash_prefill — Updates the float cache in the selected mode and applies clustered online-softmax causal prefill with the configured cache, mask, rotary, softcap, bias, or head-sink features.
• past_kv_flash_prefill_nosg — Updates the float cache in the selected mode and applies clustered online-softmax causal prefill with the configured cache, mask, rotary, softcap, bias, or head-sink features.
• new_kv_past_flash_prefill_nosg — Updates the float cache in the selected mode and applies clustered online-softmax causal prefill with the configured cache, mask, rotary, softcap, bias, or head-sink features.
• past_kv_rotary_flash_prefill_nosg — Updates the float cache in the selected mode and applies clustered online-softmax causal prefill with the configured cache, mask, rotary, softcap, bias, or head-sink features.
• past_kv_softcap_flash_prefill_nosg — Updates the float cache in the selected mode and applies clustered online-softmax causal prefill with the configured cache, mask, rotary, softcap, bias, or head-sink features.
• past_kv_headsink_flash_prefill_nosg — Updates the float cache in the selected mode and applies clustered online-softmax causal prefill with the configured cache, mask, rotary, softcap, bias, or head-sink features.
• past_kv_bias_headsink_flash_prefill_nosg — Updates the float cache in the selected mode and applies clustered online-softmax causal prefill with the configured cache, mask, rotary, softcap, bias, or head-sink features.
• window_shift_flash_prefill_nosg — Updates the float cache in the selected mode and applies clustered online-softmax causal prefill with the configured cache, mask, rotary, softcap, bias, or head-sink features.
• quant_int8_flash_prefill — Builds an int8 or int4 cache and applies clustered online-softmax prefill directly from the packed values. The family supplies subgroup and portable reductions.
• quant_int4_flash_prefill — Builds an int8 or int4 cache and applies clustered online-softmax prefill directly from the packed values. The family supplies subgroup and portable reductions.
• quant_int8_flash_prefill_nosg — Builds an int8 or int4 cache and applies clustered online-softmax prefill directly from the packed values. The family supplies subgroup and portable reductions.
• quant_int4_flash_prefill_nosg — Builds an int8 or int4 cache and applies clustered online-softmax prefill directly from the packed values. The family supplies subgroup and portable reductions.
• share_append_split_decode_splitk — Retains shared-capacity cache rows, appends new rows separately, then partitions decode across the key axis. It avoids rebuilding unchanged cache positions while exposing split-key parallelism.
• share_append_split_decode_splitk_nosg — Retains shared-capacity cache rows, appends new rows separately, then partitions decode across the key axis. It avoids rebuilding unchanged cache positions while exposing split-key parallelism.
• share_append_split_flash_prefill — Retains shared-capacity cache rows, appends new rows separately, then applies clustered causal prefill. It avoids rebuilding unchanged cache positions and includes subgroup and portable forms.
• share_append_split_flash_prefill_nosg — Retains shared-capacity cache rows, appends new rows separately, then applies clustered causal prefill. It avoids rebuilding unchanged cache positions and includes subgroup and portable forms.

Device requirements

Some implementation variants require subgroup-matrix, shader-f16, and subgroups. These are route-specific capabilities, not package-wide requirements; availability also depends on the request shape and dtype.

Files

• metadata.json — kernel metadata (id, digests, per-variant templates, provenance)
• manifest.json — the op contract (source of truth)
• test.json — correctness cases
• bench.json — benchmark + tuning cases
• attention-rank4-tiled.wgsl.jinja
• attn-flash-decode-splitk-merge.wgsl.jinja
• attn-flash-decode-splitk.wgsl.jinja
• attn-flash-online.wgsl.jinja
• attn-flash-prefill-cluster.wgsl.jinja
• attn-flash-q32-broadcast.wgsl.jinja
• attn-materialized-rowstats-combine-f32.wgsl.jinja
• attn-materialized-sgmat-f32.wgsl.jinja
• attn-online-scalar.wgsl.jinja
• gqa-attention.wgsl.jinja
• gqa-present.wgsl.jinja
• gqa-qprep.wgsl.jinja

Use with @huggingface/kernels

npm install --save-exact @huggingface/kernels@0.0.1-preview.2

Outputs with inferable metadata are allocated automatically. Explicit outputs entries request optional results or provide metadata that cannot be inferred from the supplied inputs and attributes.

This example supplies explicit metadata for:

• presentKeyT
• presentValueT

The version: 1 option selects the published kernel contract; it is independent of any operator opset, contrib since_version, or model version. It follows the v1 branch as fixes land. To pin exact artifact bytes, pass a 40-character commit revision instead of version.

Replace each *Data placeholder with a typed array containing the corresponding input data.

import { getKernel } from "@huggingface/kernels";

const kernel = await getKernel("webgpu-kernels/com.microsoft.GroupQueryAttention", { version: 1 });
// Explicit destinations request optional results or supply metadata that cannot be inferred.
const { outputT, presentKeyT, presentValueT } = await kernel({
  queryT: { data: queryTData, shape: [2, 1, 16] },
  keyT: { data: keyTData, shape: [2, 1, 8] },
  valueT: { data: valueTData, shape: [2, 1, 8] },
  pastKeyT: { data: pastKeyTData, shape: [2, 1, 8, 8] },
  pastValueT: { data: pastValueTData, shape: [2, 1, 8, 8] },
  seqlensKT: { data: seqlensKTData, shape: [2] },
  totalSequenceLengthT: { data: totalSequenceLengthTData, shape: [1] },
}, {
  attrs: { num_heads: 2, kv_num_heads: 1 },
  outputs: {
    presentKeyT: { shape: [2, 1, 8, 8], dtype: "float32" },
    presentValueT: { shape: [2, 1, 8, 8], dtype: "float32" },
  },
});