Hugging Face's logo

Kernels:
webgpu-kernels
/
com.microsoft.SparseAttention

Kernel card Files
xet
 Community

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

Name Bind key Logical dtype WebGPU storage Rank Shape Description Presence
query queryT T same as logical dtype 3 — Query (batch_size, sequence_length, num_heads * head_size), or packed [Q|K|V] (batch_size, sequence_length, (num_heads + 2 * kv_num_heads) * head_size) when key and value are omitted. required
key keyT T same as logical dtype 3 — Key (batch_size, sequence_length, kv_num_heads * head_size). Omitted for packed QKV. optional
value valueT T same as logical dtype 3 — Value (batch_size, sequence_length, kv_num_heads * head_size). Omitted for packed QKV. optional
past_key pastKeyT T same as logical dtype 4 — Key cache (batch_size, kv_num_heads, max_cache_sequence_length, head_size), updated in place. required
past_value pastValueT T same as logical dtype 4 — Value cache with the same shape as past_key, updated in place. required
block_row_indices blockRowIndicesT M int32 2 — CSR row pointers (num_layout, max_blocks + 1). Each layout starts at zero, is monotonically non-decreasing, and ends no later than that layout's block_col_indices width. required
block_col_indices blockColIndicesT M int32 2 — CSR column indices (num_layout, max_nnz_blocks), right-padded past each layout's non-zero count. Every active entry is in [0, max_blocks). required
total_sequence_length totalSequenceLengthT M int32 — — Scalar or one-element vector holding the maximum total key length. Equal to sequence_length exactly in the prompt case, which is how the past length is decided. The value fits the cache, the sparse layout's max_blocks * sparse_block_size capacity, and the rotary-cache row count when rotary is enabled. required
key_total_sequence_lengths keyTotalSequenceLengthsT M int32 1 — Per-batch total key length excluding padding, shape (batch_size). Each value is at most total_sequence_length and is at least 1 for a prompt or at least sequence_length otherwise. required
cos_cache cosCacheT T same as logical dtype 2 — Rotary cosine cache (max_rotary_sequence_length, rotary_dimension / 2), where the width is a multiple of 8 no larger than head_size / 2. Required with sin_cache when do_rotary is 1. optional
sin_cache sinCacheT T same as logical dtype 2 — Rotary sine cache with the same shape as cos_cache; required with it when do_rotary is 1. optional

Outputs

Name Bind key Logical dtype Rank Shape Description Presence
output outputT T 3 derived; see description Attention output (batch_size, sequence_length, num_heads * head_size). required
past_key pastKeyT T 4 same as past_key The key cache tensor itself after the in-place append; ONNX names this output present_key. required
past_value pastValueT T 4 same as past_value The value cache tensor itself after the in-place append; ONNX names this output present_value. required

Attributes

Attributes and default values (overridable per request):

Attribute Default Description
do_rotary 0 Set to 1 to apply rotary embedding to Q and to K before it enters the cache; every other value disables rotary embedding.
rotary_interleaved 0 Set to 1 to rotate adjacent pairs instead of using the NeoX half-split; every other value selects the NeoX layout.
num_heads — Number of query heads.
kv_num_heads — Number of key/value heads; must divide num_heads.
sparse_block_size — Tokens per sparse block; one of 16, 32, 64, 128.
scale — Scale applied to query-key products; omitted or zero uses 1 / sqrt(head_size).

Type constraints

Variable Allowed dtypes
T float32, float16
M int32

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

Use with @huggingface/kernels

The loader derives every required output's shape and logical dtype from the manifest contract and this call. It then allocates the result tensors automatically.

The version: 1 option selects the published kernel contract; it is independent of any operator opset, contrib since_version, or model 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 },
});
Downloads last month
-
kernel
webgpu
wgsl
apache-2.0
WebGPU

Requires WebGPU support. See the compatibility table.