library_name: kernels
license: apache-2.0
tags:
  - kernel
  - webgpu
  - wgsl

com.microsoft.SparseAttention

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

Description

Block-sparse causal attention used by Phi-3-small. block_row_indices and block_col_indices encode one or more CSR block masks, and layouts cycle over query heads. Grouped-query heads, separate or packed [Q|K|V], explicit scaling, partial or full rotary embedding in NeoX or interleaved layout, and float16 are supported. The past/present key and value tensors share allocations and are updated in place. Head sizes must be non-zero multiples of 8; bfloat16 is not implemented.

See the ONNX Runtime SparseAttention 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.

• separate — Portable online attention shares each key/value load across query rows. Spare lanes partition value accumulation, with the partition count bounded by the head width, workgroup size, and available workgroup storage. Single-partition staging requires one output column per lane and includes the staging array in its storage budget.
• separate_sgmat — Float32 subgroup-matrix online attention computes each score tile once, rescales output fragments through existing score scratch, and directly loads complete query tiles. Only a partial query tile is staged; key-tile width and fragment-bank capacity follow the reported workgroup-storage limit.
• separate_sgmat_tail — For a prompt, a thin query tail uses the shared portable template after complete matrix tiles. History calls retain the full matrix computation. Tail size and value partitions follow workgroup capacity. Heads narrower than its query tile use the full matrix path, since their shorter value contraction does not repay a separate portable dispatch.
• separate_rotary — Portable online attention shares each key/value load across query rows. Spare lanes partition value accumulation, with the partition count bounded by the head width, workgroup size, and available workgroup storage. Single-partition staging requires one output column per lane and includes the staging array in its storage budget.
• separate_rotary_sgmat — Float32 subgroup-matrix online attention computes each score tile once, rescales output fragments through existing score scratch, and directly loads complete query tiles. Only a partial query tile is staged; key-tile width and fragment-bank capacity follow the reported workgroup-storage limit.
• separate_rotary_sgmat_tail — For a prompt, a thin query tail uses the shared portable template after complete matrix tiles. History calls retain the full matrix computation. Tail size and value partitions follow workgroup capacity. Heads narrower than its query tile use the full matrix path, since their shorter value contraction does not repay a separate portable dispatch.
• packed — Portable online attention shares each key/value load across query rows. Spare lanes partition value accumulation, with the partition count bounded by the head width, workgroup size, and available workgroup storage. Single-partition staging requires one output column per lane and includes the staging array in its storage budget.
• packed_sgmat — Float32 subgroup-matrix online attention computes each score tile once, rescales output fragments through existing score scratch, and directly loads complete query tiles. Only a partial query tile is staged; key-tile width and fragment-bank capacity follow the reported workgroup-storage limit.
• packed_sgmat_tail — For a prompt, a thin query tail uses the shared portable template after complete matrix tiles. History calls retain the full matrix computation. Tail size and value partitions follow workgroup capacity. Heads narrower than its query tile use the full matrix path, since their shorter value contraction does not repay a separate portable dispatch.
• packed_rotary — Portable online attention shares each key/value load across query rows. Spare lanes partition value accumulation, with the partition count bounded by the head width, workgroup size, and available workgroup storage. Single-partition staging requires one output column per lane and includes the staging array in its storage budget.
• packed_rotary_sgmat — Float32 subgroup-matrix online attention computes each score tile once, rescales output fragments through existing score scratch, and directly loads complete query tiles. Only a partial query tile is staged; key-tile width and fragment-bank capacity follow the reported workgroup-storage limit.
• packed_rotary_sgmat_tail — For a prompt, a thin query tail uses the shared portable template after complete matrix tiles. History calls retain the full matrix computation. Tail size and value partitions follow workgroup capacity. Heads narrower than its query tile use the full matrix path, since their shorter value contraction does not repay a separate portable dispatch.

Device requirements

Some implementation variants require subgroup-matrix 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 cases
• sparse-attention-sgmat.wgsl.jinja
• sparse-attention.wgsl.jinja
• sparse-kv-append.wgsl.jinja
• sparse-q-rotary.wgsl.jinja

Use with @huggingface/kernels

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

Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.

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.SparseAttention", { version: 1 });
const { pastKeyT, pastValueT, outputT } = await kernel({
  queryT: { data: queryTData, shape: [1, 32, 8] },
  keyT: { data: keyTData, shape: [1, 32, 8] },
  valueT: { data: valueTData, shape: [1, 32, 8] },
  pastKeyT: { data: pastKeyTData, shape: [1, 1, 32, 8] },
  pastValueT: { data: pastValueTData, shape: [1, 1, 32, 8] },
  blockRowIndicesT: { data: blockRowIndicesTData, shape: [1, 3] },
  blockColIndicesT: { data: blockColIndicesTData, shape: [1, 3] },
  totalSequenceLengthT: { data: totalSequenceLengthTData, shape: [1] },
  keyTotalSequenceLengthsT: { data: keyTotalSequenceLengthsTData, shape: [1] },
}, {
  attrs: { num_heads: 1, kv_num_heads: 1, sparse_block_size: 16 },
});