{ "domain": "com.microsoft", "name": "MultiHeadAttention", "sinceVersion": 1, "description": "Computes multi-head self- or cross-attention over explicit query, key, and value tensors, with an optional fused QKV projection bias and additive attention bias. Supports causal masking through `unidirectional` and a configurable score scale that defaults to `1 / sqrt(head_size)`. The current kernels require query, key, and value to have the same hidden width; distinct value-head widths, KV-cache, key-padding-mask, and diagnostic-QK modes are not yet implemented.", "inputs": [ { "role": "query", "dtype": "T", "rank": 3, "description": "Query tensor of shape `(batch_size, sequence_length, hidden_size)`." }, { "role": "key", "dtype": "T", "rank": 3, "description": "Key tensor of shape `(batch_size, kv_sequence_length, hidden_size)`; the current kernels require its hidden width to equal the query hidden width." }, { "role": "value", "dtype": "T", "rank": 3, "description": "Value tensor of shape `(batch_size, kv_sequence_length, v_hidden_size)`; the current supported subset requires `v_hidden_size` to equal the query hidden width." }, { "role": "bias", "dtype": "T", "rank": 1, "optional": true, "description": "Optional fused bias from input projection with shape `(hidden_size + hidden_size + v_hidden_size)`." }, { "role": "attention_bias", "dtype": "T", "rank": 4, "optional": true, "description": "Optional additive bias applied to the attention scores before softmax, of shape `(batch_size or 1, num_heads or 1, sequence_length, kv_sequence_length)`; the first two dimensions broadcast." } ], "outputs": [ { "role": "output", "dtype": "T", "rank": 3, "shape": "[dim(shapes.queryT, 0), dim(shapes.queryT, 1), dim(shapes.valueT, 2)]", "description": "Attention output of shape `(batch_size, sequence_length, v_hidden_size)`." } ], "attributes": { "unidirectional": 0 }, "attributeDescriptions": { "num_heads": "Number of attention heads.", "scale": "Optional score scale; zero or omission selects `1 / sqrt(head_size)`.", "unidirectional": "When 1, applies an upper-left causal mask: query position `i` attends only to key positions `0..min(i, kv_sequence_length - 1)`." }, "attributeConstraints": { "num_heads": { "required": true }, "unidirectional": { "values": [0, 1] } }, "typeConstraints": { "T": ["float32", "float16"] }, "args": { "queryT": { "kind": "tensor", "semantic": "query", "role": "input" }, "keyT": { "kind": "tensor", "semantic": "key", "role": "input" }, "valueT": { "kind": "tensor", "semantic": "value", "role": "input" }, "biasT": { "kind": "tensor", "semantic": "bias", "role": "input", "required": false }, "attentionBiasT": { "kind": "tensor", "semantic": "attention_bias", "role": "input", "required": false }, "outputT": { "kind": "tensor", "semantic": "output", "role": "output" } }, "tunables": { "WORKGROUP_SIZE": 256, "SMALL_SEQ_MAX": 32, "SMALL_SEQ_MAX_PRIVATE_FLOATS": 96, "FLASH_MAX_TILE_K": 8, "FLASH_CLUSTER_WG_SMALL": 64, "FLASH_CLUSTER_WG_LARGE": 128, "FLASH_MIN_QUERY_HEADS": 248, "DECODE_MAX_SPLITS": 16, "DECODE_KEYS_PER_SPLIT": 128, "SPLITK_TARGET_WORKGROUPS": 128, "MATERIALIZED_INNER_TILE": 16, "MATERIALIZED_QUERY_TILE": 64, "MATERIALIZED_KEY_TILE": 64, "MATERIALIZED_VALUE_TILE": 64, "MATERIALIZED_VALUE_TILE_D128": 128, "MATERIALIZED_WORKGROUP_DIM": 16, "MATERIALIZED_SOFTMAX_WORKGROUP_SIZE": 256, "MATERIALIZED_SGMAT_QUERY_TILE": 64, "MATERIALIZED_SGMAT_KEY_TILE": 64, "MATERIALIZED_SGMAT_INNER_TILE": 32, "MATERIALIZED_CACHED_SOFTMAX_WORKGROUP_SIZE": 128, "MATERIALIZED_CACHED_SOFTMAX_MAX_VECS_PER_LANE": 4, "PREFILL_QUERY_TILE": 32, "MATERIALIZED_FUSED_SOFTMAX_MIN_SCORE_BYTES": 16777216 }, "tunableDescriptions": { "SMALL_SEQ_MAX": "Maximum query and key length admitted by the whole-head small-sequence route. It is a resource ceiling; flash eligibility determines the crossover.", "SMALL_SEQ_MAX_PRIVATE_FLOATS": "Caps the combined per-lane query-row and score private arrays for the small-sequence route so corner shapes stay within its register budget.", "FLASH_MIN_QUERY_HEADS": "Minimum query-row/head work count used to admit flash attention unless a long key axis admits it independently.", "DECODE_MAX_SPLITS": "Caps the number of key-axis partitions created by split-K attention.", "DECODE_KEYS_PER_SPLIT": "Targets this many cached keys per partition when deriving the split-K count.", "SPLITK_TARGET_WORKGROUPS": "Targets total workgroups for short-query split-K; the partition count grows until independent query/head work approaches this value.", "MATERIALIZED_CACHED_SOFTMAX_MAX_VECS_PER_LANE": "Caps cached score vectors owned by each softmax lane; exceeding it rejects the cached materialized-softmax route.", "MATERIALIZED_FUSED_SOFTMAX_MIN_SCORE_BYTES": "Minimum score-matrix size for fusing row statistics into materialization, avoiding a full softmax read and write when that tradeoff is worthwhile." }, "derive": { "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)", "wave32Adapter": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize == 32 and device.adapterInfo.subgroupMaxSize == 32", "subgroupsWave32": "device.features.has(\"subgroups\") and wave32Adapter", "narrowSubgroupRange": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize < device.adapterInfo.subgroupMaxSize and device.adapterInfo.subgroupMaxSize <= 16", "canPinSubgroupSize32": "device.features.has(\"subgroups\") and device.features.has(\"subgroup-size-control\") and has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize <= 32 and device.adapterInfo.subgroupMaxSize >= 32", "pinSubgroupSize32": "canPinSubgroupSize32 and not wave32Adapter", "wave32Effective": "wave32Adapter or pinSubgroupSize32", "headDim": "dim(shapes.query, 2) / attrs.num_heads if (ranks.query == 3 and attrs.num_heads > 0) else 0", "qkvDtypesOk": "tensorDtypes.key == tensorDtypes.query and tensorDtypes.value == tensorDtypes.query and tensorDtypes.output == tensorDtypes.query", "floatDtypeOk": "(tensorDtypes.query == \"float32\" or tensorDtypes.query == \"float16\") and f16Ok(tensorDtypes.query)", "qkvShapeOk": "ranks.query == 3 and ranks.key == 3 and ranks.value == 3 and ranks.output == 3 and attrs.num_heads > 0 and dim(shapes.query, 2) % attrs.num_heads == 0 and dim(shapes.key, 2) == dim(shapes.query, 2) and dim(shapes.value, 2) == dim(shapes.query, 2) and dim(shapes.key, 1) == dim(shapes.value, 1) and dim(shapes.query, 0) == dim(shapes.key, 0) and dim(shapes.query, 0) == dim(shapes.value, 0) and dim(shapes.output, 0) == dim(shapes.query, 0) and dim(shapes.output, 1) == dim(shapes.query, 1) and dim(shapes.output, 2) == dim(shapes.value, 2)", "noAttnBias": "not present.attentionBiasT", "attnBiasOk": "present.attentionBiasT and ranks.attention_bias == 4 and tensorDtypes.attention_bias == tensorDtypes.query and (dim(shapes.attention_bias, 0) == dim(shapes.query, 0) or dim(shapes.attention_bias, 0) == 1) and (dim(shapes.attention_bias, 1) == attrs.num_heads or dim(shapes.attention_bias, 1) == 1) and dim(shapes.attention_bias, 2) == dim(shapes.query, 1) and dim(shapes.attention_bias, 3) == dim(shapes.key, 1)", "qkvContractOk": "qkvShapeOk and qkvDtypesOk and floatDtypeOk and noAttnBias", "qkvMaskContractOk": "qkvShapeOk and qkvDtypesOk and floatDtypeOk and attnBiasOk", "biasOk": "present.biasT and ranks.bias == 1 and tensorDtypes.bias == tensorDtypes.query and dim(shapes.bias, 0) == 3 * dim(shapes.query, 2)", "q32BroadcastSubgroupLanes": "device.adapterInfo.subgroupMinSize if subgroupsWave32 else 0", "q32BroadcastF32HeadVectors": "headDim / 4 if headDim % 4 == 0 else 0", "q32BroadcastF32RegisterGeometry": "subgroupsWave32 and q32BroadcastF32HeadVectors == q32BroadcastSubgroupLanes", "subgroupCluster4": "not narrowSubgroupRange and device.features.has(\"subgroups\") and has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize >= 4 and device.adapterInfo.subgroupMinSize % 4 == 0 and device.adapterInfo.subgroupMaxSize % 4 == 0", "subgroupCluster8": "not narrowSubgroupRange and device.features.has(\"subgroups\") and has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize >= 8 and device.adapterInfo.subgroupMinSize % 8 == 0 and device.adapterInfo.subgroupMaxSize % 8 == 0", "attentionDispatchFits": "dim(shapes.query, 0) <= device.limits.maxComputeWorkgroupsPerDimension and dim(shapes.query, 1) <= device.limits.maxComputeWorkgroupsPerDimension and attrs.num_heads <= device.limits.maxComputeWorkgroupsPerDimension", "flashHeadOk": "headDim % 4 == 0 and headDim >= 32 and headDim <= 256", "flashSizeOk": "flashHeadOk and (dim(shapes.query, 1) * attrs.num_heads >= tunables.FLASH_MIN_QUERY_HEADS or (dim(shapes.query, 1) == 1 and dim(shapes.key, 1) >= 512) or (dim(shapes.query, 1) > 1 and dim(shapes.key, 1) >= 2048)) and attentionDispatchFits", "flashShapeOk": "qkvContractOk and flashSizeOk", "flashMaskShapeOk": "qkvMaskContractOk and flashSizeOk", "noBiasSplitKCount": "min(tunables.DECODE_MAX_SPLITS if dim(shapes.query, 1) == 1 else max(1, ceilDiv(tunables.SPLITK_TARGET_WORKGROUPS, dim(shapes.query, 0) * dim(shapes.query, 1) * attrs.num_heads)), ceilDiv(dim(shapes.key, 1), tunables.DECODE_KEYS_PER_SPLIT))", "biasSplitKCount": "min(tunables.DECODE_MAX_SPLITS, ceilDiv(dim(shapes.key, 1), tunables.DECODE_KEYS_PER_SPLIT))", "noBiasPartialOutBytes": "dim(shapes.query, 0) * dim(shapes.query, 1) * attrs.num_heads * noBiasSplitKCount * headDim * 4", "noBiasStatsBytes": "2 * dim(shapes.query, 0) * dim(shapes.query, 1) * attrs.num_heads * noBiasSplitKCount * 4", "noBiasSplitScratchFits": "noBiasPartialOutBytes <= device.limits.maxStorageBufferBindingSize and noBiasPartialOutBytes <= device.limits.maxBufferSize and noBiasStatsBytes <= device.limits.maxStorageBufferBindingSize and noBiasStatsBytes <= device.limits.maxBufferSize", "noBiasSplitDispatchFits": "dim(shapes.query, 1) * noBiasSplitKCount <= device.limits.maxComputeWorkgroupsPerDimension", "biasPartialOutBytes": "dim(shapes.query, 0) * attrs.num_heads * biasSplitKCount * headDim * 4", "biasStatsBytes": "2 * dim(shapes.query, 0) * attrs.num_heads * biasSplitKCount * 4", "biasSplitScratchFits": "biasPartialOutBytes <= device.limits.maxStorageBufferBindingSize and biasPartialOutBytes <= device.limits.maxBufferSize and biasStatsBytes <= device.limits.maxStorageBufferBindingSize and biasStatsBytes <= device.limits.maxBufferSize", "biasSplitDispatchFits": "biasSplitKCount <= device.limits.maxComputeWorkgroupsPerDimension", "decodeSplitKNoBiasOk": "qkvContractOk and attrs.unidirectional == 0 and dim(shapes.query, 1) == 1 and dim(shapes.key, 1) >= 512 and flashHeadOk and attentionDispatchFits and noBiasSplitDispatchFits and noBiasSplitScratchFits", "shortQuerySplitKNoBiasOk": "qkvContractOk and attrs.unidirectional == 0 and dim(shapes.query, 1) >= 2 and dim(shapes.query, 1) <= 16 and dim(shapes.key, 1) >= 2048 and flashHeadOk and attentionDispatchFits and noBiasSplitDispatchFits and noBiasSplitScratchFits", "decodeSplitKBiasOk": "biasOk and qkvContractOk and attrs.unidirectional == 0 and dim(shapes.query, 1) == 1 and dim(shapes.key, 1) >= 512 and flashHeadOk and attentionDispatchFits and biasSplitDispatchFits and biasSplitScratchFits", "decodeSplitKPortablePreferred": "tensorDtypes.query == \"float32\" and device.features.has(\"subgroups\") and has(device.adapterInfo, \"subgroupMinSize\") and headDim / 4 < device.adapterInfo.subgroupMinSize", "materializedScoreBytes": "dim(shapes.query, 0) * attrs.num_heads * dim(shapes.query, 1) * dim(shapes.key, 1) * 4", "materializedScoreFits": "materializedScoreBytes <= device.limits.maxStorageBufferBindingSize and materializedScoreBytes <= device.limits.maxBufferSize", "materializedWorkgroupSize": "tunables.MATERIALIZED_WORKGROUP_DIM * tunables.MATERIALIZED_WORKGROUP_DIM", "materializedScoreStorageBytes": "(tunables.MATERIALIZED_QUERY_TILE * tunables.MATERIALIZED_INNER_TILE + tunables.MATERIALIZED_KEY_TILE * (tunables.MATERIALIZED_INNER_TILE + 4)) * 4", "materializedApplyTileN": "tunables.MATERIALIZED_VALUE_TILE_D128 if headDim == 128 else tunables.MATERIALIZED_VALUE_TILE", "materializedApplyStorageBytes": "(tunables.MATERIALIZED_QUERY_TILE * tunables.MATERIALIZED_INNER_TILE + tunables.MATERIALIZED_INNER_TILE * materializedApplyTileN) * 4", "materializedTileGeometryOk": "tunables.MATERIALIZED_INNER_TILE % 4 == 0 and (materializedApplyTileN / tunables.MATERIALIZED_WORKGROUP_DIM) % 4 == 0 and tunables.MATERIALIZED_INNER_TILE > 0 and tunables.MATERIALIZED_WORKGROUP_DIM > 0 and tunables.MATERIALIZED_QUERY_TILE % tunables.MATERIALIZED_WORKGROUP_DIM == 0 and tunables.MATERIALIZED_KEY_TILE % tunables.MATERIALIZED_WORKGROUP_DIM == 0 and materializedApplyTileN % tunables.MATERIALIZED_WORKGROUP_DIM == 0 and tunables.MATERIALIZED_QUERY_TILE * tunables.MATERIALIZED_INNER_TILE >= materializedWorkgroupSize and tunables.MATERIALIZED_KEY_TILE * tunables.MATERIALIZED_INNER_TILE >= materializedWorkgroupSize and tunables.MATERIALIZED_INNER_TILE * materializedApplyTileN >= materializedWorkgroupSize and (tunables.MATERIALIZED_QUERY_TILE * tunables.MATERIALIZED_INNER_TILE) % materializedWorkgroupSize == 0 and (tunables.MATERIALIZED_KEY_TILE * tunables.MATERIALIZED_INNER_TILE) % materializedWorkgroupSize == 0 and (tunables.MATERIALIZED_INNER_TILE * materializedApplyTileN) % materializedWorkgroupSize == 0", "materializedDeviceOk": "materializedTileGeometryOk and materializedWorkgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup and tunables.MATERIALIZED_WORKGROUP_DIM <= device.limits.maxComputeWorkgroupSizeX and tunables.MATERIALIZED_WORKGROUP_DIM <= device.limits.maxComputeWorkgroupSizeY and materializedScoreStorageBytes <= device.limits.maxComputeWorkgroupStorageSize and materializedApplyStorageBytes <= device.limits.maxComputeWorkgroupStorageSize", "materializedWideSimdOk": "device.features.has(\"subgroups\") or (has(device.adapterInfo, \"subgroupMinSize\") and device.adapterInfo.subgroupMinSize >= 16)", "materializedF32CoreOk": "qkvContractOk and tensorDtypes.query == \"float32\" and attrs.unidirectional == 0 and headDim >= 64 and headDim <= 128 and dim(shapes.query, 1) >= 512 and dim(shapes.key, 1) >= 512 and dim(shapes.query, 0) * attrs.num_heads <= device.limits.maxComputeWorkgroupsPerDimension and ceilDiv(dim(shapes.query, 1), tunables.MATERIALIZED_QUERY_TILE) <= device.limits.maxComputeWorkgroupsPerDimension and ceilDiv(dim(shapes.key, 1), tunables.MATERIALIZED_KEY_TILE) <= device.limits.maxComputeWorkgroupsPerDimension and materializedScoreFits and materializedDeviceOk", "materializedSoftmaxStorageBytes": "tunables.MATERIALIZED_SOFTMAX_WORKGROUP_SIZE * 8 + 8", "materializedSoftmaxResourcesFit": "tunables.MATERIALIZED_SOFTMAX_WORKGROUP_SIZE > 0 and pow2ceil(tunables.MATERIALIZED_SOFTMAX_WORKGROUP_SIZE) == tunables.MATERIALIZED_SOFTMAX_WORKGROUP_SIZE and tunables.MATERIALIZED_SOFTMAX_WORKGROUP_SIZE <= deviceWorkgroupCap and materializedSoftmaxStorageBytes <= device.limits.maxComputeWorkgroupStorageSize", "materializedSgmatQueryTile": "tunables.MATERIALIZED_SGMAT_QUERY_TILE", "materializedSgmatKeyTile": "tunables.MATERIALIZED_SGMAT_KEY_TILE", "materializedSgmatInnerTile": "tunables.MATERIALIZED_SGMAT_INNER_TILE", "materializedSgmatSubgroupRows": "floor(materializedSgmatQueryTile / 16)", "materializedSgmatSubgroupCols": "floor(materializedSgmatKeyTile / 32)", "materializedSgmatStatSlots": "ceilDiv(dim(shapes.key, 1), materializedSgmatKeyTile) * materializedSgmatSubgroupCols", "materializedRowStatsWg": "min(tunables.MATERIALIZED_SOFTMAX_WORKGROUP_SIZE, deviceWorkgroupCap)", "materializedRowStatsElements": "dim(shapes.query, 0) * attrs.num_heads * dim(shapes.query, 1) * 2", "materializedScorePartialElements": "dim(shapes.query, 0) * attrs.num_heads * materializedSgmatStatSlots * dim(shapes.query, 1) * 2", "materializedGemmStatSlots": "ceilDiv(dim(shapes.key, 1), tunables.MATERIALIZED_KEY_TILE)", "materializedGemmScorePartialElements": "dim(shapes.query, 0) * attrs.num_heads * materializedGemmStatSlots * dim(shapes.query, 1) * 2", "materializedGemmFusedSoftmaxWorthIt": "materializedScoreBytes >= tunables.MATERIALIZED_FUSED_SOFTMAX_MIN_SCORE_BYTES and materializedGemmScorePartialElements * 4 <= device.limits.maxStorageBufferBindingSize and materializedGemmScorePartialElements * 4 <= device.limits.maxBufferSize and materializedRowStatsElements * 4 <= device.limits.maxStorageBufferBindingSize and materializedRowStatsWg > 0 and tunables.MATERIALIZED_QUERY_TILE <= tunables.MATERIALIZED_WORKGROUP_DIM * tunables.MATERIALIZED_WORKGROUP_DIM", "materializedFusedSoftmaxTilesOk": "materializedSgmatQueryTile >= 64 and materializedSgmatKeyTile >= 64", "materializedFusedSoftmaxWorthIt": "materializedScoreBytes >= tunables.MATERIALIZED_FUSED_SOFTMAX_MIN_SCORE_BYTES and materializedFusedSoftmaxTilesOk", "materializedSgmatWorkgroupSize": "materializedSgmatSubgroupRows * materializedSgmatSubgroupCols * 32", "materializedSgmatCompactStorageBytes": "(materializedSgmatQueryTile + materializedSgmatKeyTile) * materializedSgmatInnerTile * 4", "materializedSgmatGeometryOk": "materializedSgmatQueryTile >= 16 and materializedSgmatQueryTile % 16 == 0 and materializedSgmatKeyTile >= 32 and materializedSgmatKeyTile <= 64 and materializedSgmatKeyTile % 32 == 0 and materializedSgmatInnerTile == 32 and materializedSgmatWorkgroupSize > 0", "materializedSgmatBuffersFit": "numel(shapes.query) * 4 <= device.limits.maxStorageBufferBindingSize and numel(shapes.query) * 4 <= device.limits.maxBufferSize and numel(shapes.key) * 4 <= device.limits.maxStorageBufferBindingSize and numel(shapes.key) * 4 <= device.limits.maxBufferSize and numel(shapes.value) * 4 <= device.limits.maxStorageBufferBindingSize and numel(shapes.value) * 4 <= device.limits.maxBufferSize and numel(shapes.output) * 4 <= device.limits.maxStorageBufferBindingSize and numel(shapes.output) * 4 <= device.limits.maxBufferSize and materializedScoreFits", "materializedSgmatResourcesFit": "materializedSgmatGeometryOk and materializedSgmatWorkgroupSize <= deviceWorkgroupCap and materializedSgmatCompactStorageBytes <= device.limits.maxComputeWorkgroupStorageSize", "materializedSgmatDispatchFits": "ceilDiv(dim(shapes.key, 1), materializedSgmatKeyTile) <= device.limits.maxComputeWorkgroupsPerDimension and ceilDiv(dim(shapes.query, 1), materializedSgmatQueryTile) <= device.limits.maxComputeWorkgroupsPerDimension and dim(shapes.query, 0) * attrs.num_heads <= device.limits.maxComputeWorkgroupsPerDimension and ceilDiv(dim(shapes.query, 0) * attrs.num_heads * dim(shapes.query, 1), device.limits.maxComputeWorkgroupsPerDimension) <= device.limits.maxComputeWorkgroupsPerDimension", "materializedSgmatDirectScoreStore": "dim(shapes.query, 1) % materializedSgmatQueryTile == 0 and dim(shapes.key, 1) % materializedSgmatKeyTile == 0", "materializedSgmatDirectApplyStore": "dim(shapes.query, 1) % materializedSgmatQueryTile == 0 and headDim % materializedSgmatKeyTile == 0", "materializedSgmatRuntimeDirectStore": "dim(shapes.query, 1) >= 2 * materializedSgmatQueryTile and dim(shapes.key, 1) >= 2 * materializedSgmatKeyTile", "materializedSgmatCoreOk": "qkvContractOk and tensorDtypes.query == \"float32\" and attrs.unidirectional == 0 and headDim >= 32 and headDim <= 256 and dim(shapes.query, 1) >= 512 and dim(shapes.key, 1) >= 512 and device.features.has(\"subgroups\") and wave32Effective and device.features.has(\"chromium-experimental-subgroup-matrix\") and materializedSgmatBuffersFit and materializedSgmatResourcesFit and materializedSgmatDispatchFits", "materializedCachedSoftmaxWg": "min(tunables.MATERIALIZED_CACHED_SOFTMAX_WORKGROUP_SIZE, device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)", "materializedCachedSoftmaxVecsPerLane": "ceilDiv(ceilDiv(dim(shapes.key, 1), 4), max(1, materializedCachedSoftmaxWg))", "materializedCachedSoftmaxStorageBytes": "materializedCachedSoftmaxWg * 8 + 8", "materializedCachedSoftmaxResourcesFit": "materializedCachedSoftmaxWg > 0 and pow2ceil(materializedCachedSoftmaxWg) == materializedCachedSoftmaxWg and materializedCachedSoftmaxStorageBytes <= device.limits.maxComputeWorkgroupStorageSize and tunables.MATERIALIZED_CACHED_SOFTMAX_MAX_VECS_PER_LANE > 0 and materializedCachedSoftmaxVecsPerLane <= tunables.MATERIALIZED_CACHED_SOFTMAX_MAX_VECS_PER_LANE", "materializedCachedSoftmaxOk": "dim(shapes.key, 1) % 4 == 0 and materializedCachedSoftmaxResourcesFit", "materializedAdaptiveSoftmaxOk": "materializedCachedSoftmaxOk or materializedSoftmaxResourcesFit", "materializedSgmatOk": "materializedSgmatCoreOk and materializedAdaptiveSoftmaxOk", "materializedSgmatFusedOk": "materializedSgmatCoreOk", "materializedF32Ok": "materializedF32CoreOk and materializedAdaptiveSoftmaxOk", "clusterTileKWg64": "max(1, min(tunables.FLASH_MAX_TILE_K, floor(device.limits.maxComputeWorkgroupStorageSize / (headDim * (8 if tensorDtypes.query != \"float16\" else 4) + tunables.FLASH_CLUSTER_WG_SMALL * 4))))", "clusterTileKWg128": "max(1, min(tunables.FLASH_MAX_TILE_K, floor(device.limits.maxComputeWorkgroupStorageSize / (headDim * (8 if tensorDtypes.query != \"float16\" else 4) + tunables.FLASH_CLUSTER_WG_LARGE * 4))))", "smallHeadParallelOk": "qkvContractOk and headDim < 32 and dim(shapes.key, 1) >= 64 and dim(shapes.key, 1) <= 2048 and attrs.unidirectional == 0 and dim(shapes.query, 0) <= device.limits.maxComputeWorkgroupsPerDimension and attrs.num_heads <= device.limits.maxComputeWorkgroupsPerDimension", "prefillTiledStorageBytes": "headDim * tunables.PREFILL_QUERY_TILE * 4", "prefillTiledDeviceOk": "tunables.PREFILL_QUERY_TILE <= deviceWorkgroupCap and prefillTiledStorageBytes <= device.limits.maxComputeWorkgroupStorageSize", "portableWorkgroupSize": "min(tunables.WORKGROUP_SIZE, pow2ceil(max(1, headDim)))", "portableWorkgroupStorageBytes": "portableWorkgroupSize * 4 + max(1, headDim) * 4 + 16", "portableWorkgroupOk": "tunables.WORKGROUP_SIZE > 0 and pow2ceil(tunables.WORKGROUP_SIZE) == tunables.WORKGROUP_SIZE and portableWorkgroupSize <= deviceWorkgroupCap and portableWorkgroupStorageBytes <= device.limits.maxComputeWorkgroupStorageSize", "fallbackShapeOk": "qkvContractOk and portableWorkgroupOk", "fallbackMaskShapeOk": "qkvMaskContractOk and portableWorkgroupOk", "smallSeqShapeOk": "qkvContractOk and tensorDtypes.query == \"float32\" and dim(shapes.query, 1) >= 1 and dim(shapes.query, 1) <= tunables.SMALL_SEQ_MAX and dim(shapes.key, 1) >= 1 and dim(shapes.key, 1) <= tunables.SMALL_SEQ_MAX and headDim >= 1 and attrs.unidirectional == 0", "smallSeqPrivateFloats": "dim(shapes.key, 1) + headDim", "smallSeqWorkgroupSize": "max(32, pow2ceil(dim(shapes.query, 1)))", "smallSeqSharedBytes": "dim(shapes.key, 1) * headDim * 8", "smallSeqResourcesFit": "smallSeqPrivateFloats <= tunables.SMALL_SEQ_MAX_PRIVATE_FLOATS and smallSeqWorkgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup and smallSeqWorkgroupSize <= device.limits.maxComputeWorkgroupSizeX and smallSeqSharedBytes <= device.limits.maxComputeWorkgroupStorageSize", "smallSeqDispatchFits": "attrs.num_heads <= device.limits.maxComputeWorkgroupsPerDimension and dim(shapes.query, 0) <= device.limits.maxComputeWorkgroupsPerDimension", "materializedSgmatCoreF16Ok": "qkvContractOk and tensorDtypes.query == \"float16\" and tensorDtypes.key == \"float16\" and tensorDtypes.value == \"float16\" and device.features.has(\"shader-f16\") and attrs.unidirectional == 0 and headDim >= 32 and headDim <= 256 and dim(shapes.query, 1) >= 512 and dim(shapes.key, 1) >= 512 and device.features.has(\"subgroups\") and wave32Effective and device.features.has(\"chromium-experimental-subgroup-matrix\") and materializedSgmatBuffersFit and materializedSgmatResourcesFit and materializedSgmatDispatchFits", "materializedSgmatFusedF16Ok": "materializedSgmatCoreF16Ok" }, "constants": { "attentionScaleExpression": "\"select(inverseSqrt(f32(HEAD_DIM)), params.scale, params.scale != 0.0)\"" }, "bindingSets": { "decodeSplitPartial": [ { "name": "query", "arg": "queryT", "semantic": "query", "buffer": { "type": "read-only-storage" }, "elementType": "$inputVec4" }, { "name": "key", "arg": "keyT", "semantic": "key", "buffer": { "type": "read-only-storage" }, "elementType": "$inputVec4" }, { "name": "value", "arg": "valueT", "semantic": "value", "buffer": { "type": "read-only-storage" }, "elementType": "$inputVec4" }, { "name": "partial_out", "semantic": "partialOut", "buffer": { "type": "storage" }, "elementType": "vec4" }, { "name": "partial_stats", "semantic": "partialStats", "buffer": { "type": "storage" }, "elementType": "vec2" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [ { "name": "kvSeq", "type": "u32", "value": "dim(shapes.key, 1)" }, { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" } ] } } ], "decodeSplitMerge": [ { "name": "partial_out", "semantic": "partialOut", "buffer": { "type": "read-only-storage" }, "elementType": "vec4" }, { "name": "partial_stats", "semantic": "partialStats", "buffer": { "type": "read-only-storage" }, "elementType": "vec2" }, { "name": "output", "arg": "outputT", "semantic": "output", "buffer": { "type": "storage" }, "elementType": "$inputVec4" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [ { "name": "qSeq", "type": "u32", "value": "dim(shapes.query, 1)" }, { "name": "kvSeq", "type": "u32", "value": "dim(shapes.key, 1)" }, { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" } ] } } ], "decodeSplitMergeNoParams": [ { "name": "partial_out", "semantic": "partialOut", "buffer": { "type": "read-only-storage" }, "elementType": "vec4" }, { "name": "partial_stats", "semantic": "partialStats", "buffer": { "type": "read-only-storage" }, "elementType": "vec2" }, { "name": "output", "arg": "outputT", "semantic": "output", "buffer": { "type": "storage" }, "elementType": "$inputVec4" } ], "noBias": [ { "name": "query", "arg": "queryT", "semantic": "query", "buffer": { "type": "read-only-storage" }, "elementType": "$inputElement" }, { "name": "key", "arg": "keyT", "semantic": "key", "buffer": { "type": "read-only-storage" }, "elementType": "$inputElement" }, { "name": "value", "arg": "valueT", "semantic": "value", "buffer": { "type": "read-only-storage" }, "elementType": "$inputElement" }, { "name": "output", "arg": "outputT", "semantic": "output", "buffer": { "type": "storage" }, "elementType": "$outputElement" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [ { "name": "qSeq", "type": "u32", "value": "dim(shapes.query, 1)" }, { "name": "kvSeq", "type": "u32", "value": "dim(shapes.key, 1)" }, { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" }, { "name": "isCausal", "type": "u32", "value": "attrs.unidirectional" } ] } } ], "noBiasSmallBidirectional": [ { "name": "query", "arg": "queryT", "semantic": "query", "buffer": { "type": "read-only-storage" }, "elementType": "$inputElement" }, { "name": "key", "arg": "keyT", "semantic": "key", "buffer": { "type": "read-only-storage" }, "elementType": "$inputElement" }, { "name": "value", "arg": "valueT", "semantic": "value", "buffer": { "type": "read-only-storage" }, "elementType": "$inputElement" }, { "name": "output", "arg": "outputT", "semantic": "output", "buffer": { "type": "storage" }, "elementType": "$outputElement" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [ { "name": "qSeq", "type": "u32", "value": "dim(shapes.query, 1)" }, { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" } ] } } ], "bias": [ { "name": "query", "arg": "queryT", "semantic": "query", "buffer": { "type": "read-only-storage" }, "elementType": "$inputElement" }, { "name": "key", "arg": "keyT", "semantic": "key", "buffer": { "type": "read-only-storage" }, "elementType": "$inputElement" }, { "name": "value", "arg": "valueT", "semantic": "value", "buffer": { "type": "read-only-storage" }, "elementType": "$inputElement" }, { "name": "bias", "arg": "biasT", "semantic": "bias", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, { "name": "output", "arg": "outputT", "semantic": "output", "buffer": { "type": "storage" }, "elementType": "$outputElement" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [ { "name": "qSeq", "type": "u32", "value": "dim(shapes.query, 1)" }, { "name": "kvSeq", "type": "u32", "value": "dim(shapes.key, 1)" }, { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" }, { "name": "isCausal", "type": "u32", "value": "attrs.unidirectional" } ] } } ], "biasNonCausal": [ { "name": "query", "arg": "queryT", "semantic": "query", "buffer": { "type": "read-only-storage" }, "elementType": "$inputElement" }, { "name": "key", "arg": "keyT", "semantic": "key", "buffer": { "type": "read-only-storage" }, "elementType": "$inputElement" }, { "name": "value", "arg": "valueT", "semantic": "value", "buffer": { "type": "read-only-storage" }, "elementType": "$inputElement" }, { "name": "bias", "arg": "biasT", "semantic": "bias", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, { "name": "output", "arg": "outputT", "semantic": "output", "buffer": { "type": "storage" }, "elementType": "$outputElement" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [ { "name": "qSeq", "type": "u32", "value": "dim(shapes.query, 1)" }, { "name": "kvSeq", "type": "u32", "value": "dim(shapes.key, 1)" }, { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" } ] } } ], "noBiasMask": [ { "name": "query", "arg": "queryT", "semantic": "query", "buffer": { "type": "read-only-storage" }, "elementType": "$inputElement" }, { "name": "key", "arg": "keyT", "semantic": "key", "buffer": { "type": "read-only-storage" }, "elementType": "$inputElement" }, { "name": "value", "arg": "valueT", "semantic": "value", "buffer": { "type": "read-only-storage" }, "elementType": "$inputElement" }, { "name": "attn_mask", "arg": "attentionBiasT", "semantic": "attention_bias", "buffer": { "type": "read-only-storage" }, "elementType": "$maskElement" }, { "name": "output", "arg": "outputT", "semantic": "output", "buffer": { "type": "storage" }, "elementType": "$outputElement" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [ { "name": "qSeq", "type": "u32", "value": "dim(shapes.query, 1)" }, { "name": "kvSeq", "type": "u32", "value": "dim(shapes.key, 1)" }, { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" }, { "name": "isCausal", "type": "u32", "value": "attrs.unidirectional" }, { "name": "maskBatchStride", "type": "u32", "value": "0 if dim(shapes.attention_bias, 0) == 1 else dim(shapes.attention_bias, 1) * dim(shapes.attention_bias, 2) * dim(shapes.attention_bias, 3)" }, { "name": "maskHeadStride", "type": "u32", "value": "0 if dim(shapes.attention_bias, 1) == 1 else dim(shapes.attention_bias, 2) * dim(shapes.attention_bias, 3)" }, { "name": "maskSeqStride", "type": "u32", "value": "dim(shapes.attention_bias, 3)" } ] } } ], "biasMask": [ { "name": "query", "arg": "queryT", "semantic": "query", "buffer": { "type": "read-only-storage" }, "elementType": "$inputElement" }, { "name": "key", "arg": "keyT", "semantic": "key", "buffer": { "type": "read-only-storage" }, "elementType": "$inputElement" }, { "name": "value", "arg": "valueT", "semantic": "value", "buffer": { "type": "read-only-storage" }, "elementType": "$inputElement" }, { "name": "attn_mask", "arg": "attentionBiasT", "semantic": "attention_bias", "buffer": { "type": "read-only-storage" }, "elementType": "$maskElement" }, { "name": "bias", "arg": "biasT", "semantic": "bias", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, { "name": "output", "arg": "outputT", "semantic": "output", "buffer": { "type": "storage" }, "elementType": "$outputElement" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [ { "name": "qSeq", "type": "u32", "value": "dim(shapes.query, 1)" }, { "name": "kvSeq", "type": "u32", "value": "dim(shapes.key, 1)" }, { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" }, { "name": "isCausal", "type": "u32", "value": "attrs.unidirectional" }, { "name": "maskBatchStride", "type": "u32", "value": "0 if dim(shapes.attention_bias, 0) == 1 else dim(shapes.attention_bias, 1) * dim(shapes.attention_bias, 2) * dim(shapes.attention_bias, 3)" }, { "name": "maskHeadStride", "type": "u32", "value": "0 if dim(shapes.attention_bias, 1) == 1 else dim(shapes.attention_bias, 2) * dim(shapes.attention_bias, 3)" }, { "name": "maskSeqStride", "type": "u32", "value": "dim(shapes.attention_bias, 3)" } ] } } ], "decodeSplitBiasPartial": [ { "name": "query", "arg": "queryT", "semantic": "query", "buffer": { "type": "read-only-storage" }, "elementType": "$inputVec4" }, { "name": "key", "arg": "keyT", "semantic": "key", "buffer": { "type": "read-only-storage" }, "elementType": "$inputVec4" }, { "name": "value", "arg": "valueT", "semantic": "value", "buffer": { "type": "read-only-storage" }, "elementType": "$inputVec4" }, { "name": "bias", "arg": "biasT", "semantic": "bias", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, { "name": "partial_out", "semantic": "partialOut", "buffer": { "type": "storage" }, "elementType": "vec4" }, { "name": "partial_stats", "semantic": "partialStats", "buffer": { "type": "storage" }, "elementType": "vec2" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [ { "name": "kvSeq", "type": "u32", "value": "dim(shapes.key, 1)" }, { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" } ] } } ], "decodeSplitBiasMerge": [ { "name": "partial_out", "semantic": "partialOut", "buffer": { "type": "read-only-storage" }, "elementType": "vec4" }, { "name": "partial_stats", "semantic": "partialStats", "buffer": { "type": "read-only-storage" }, "elementType": "vec2" }, { "name": "bias", "arg": "biasT", "semantic": "bias", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, { "name": "output", "arg": "outputT", "semantic": "output", "buffer": { "type": "storage" }, "elementType": "$inputVec4" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [ { "name": "qSeq", "type": "u32", "value": "dim(shapes.query, 1)" }, { "name": "kvSeq", "type": "u32", "value": "dim(shapes.key, 1)" }, { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" } ] } } ], "decodeSplitBiasMergeNoParams": [ { "name": "partial_out", "semantic": "partialOut", "buffer": { "type": "read-only-storage" }, "elementType": "vec4" }, { "name": "partial_stats", "semantic": "partialStats", "buffer": { "type": "read-only-storage" }, "elementType": "vec2" }, { "name": "bias", "arg": "biasT", "semantic": "bias", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, { "name": "output", "arg": "outputT", "semantic": "output", "buffer": { "type": "storage" }, "elementType": "$inputVec4" } ], "materializedScore": [ { "name": "query", "arg": "queryT", "semantic": "query", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, { "name": "key", "arg": "keyT", "semantic": "key", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, { "name": "scores", "semantic": "materializedScores", "buffer": { "type": "storage" }, "elementType": "f32" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [ { "name": "qSeq", "type": "u32", "value": "dim(shapes.query, 1)" }, { "name": "kvSeq", "type": "u32", "value": "dim(shapes.key, 1)" }, { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" } ] } } ], "materializedScoreBias": [ { "name": "query", "arg": "queryT", "semantic": "query", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, { "name": "key", "arg": "keyT", "semantic": "key", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, { "name": "bias", "arg": "biasT", "semantic": "bias", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, { "name": "scores", "semantic": "materializedScores", "buffer": { "type": "storage" }, "elementType": "f32" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [ { "name": "qSeq", "type": "u32", "value": "dim(shapes.query, 1)" }, { "name": "kvSeq", "type": "u32", "value": "dim(shapes.key, 1)" }, { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" } ] } } ], "materializedSoftmax": [ { "name": "scores", "semantic": "materializedScores", "buffer": { "type": "storage" }, "elementType": "$softmaxElementType" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [ { "name": "rows", "type": "u32", "value": "dim(shapes.query, 0) * attrs.num_heads * dim(shapes.query, 1)" } ] } } ], "materializedApply": [ { "name": "scores", "semantic": "materializedScores", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, { "name": "value", "arg": "valueT", "semantic": "value", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, { "name": "output", "arg": "outputT", "semantic": "output", "buffer": { "type": "storage" }, "elementType": "f32" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [ { "name": "qSeq", "type": "u32", "value": "dim(shapes.query, 1)" }, { "name": "kvSeq", "type": "u32", "value": "dim(shapes.key, 1)" } ] } } ], "prefillTiledNoBias": [ { "name": "q", "arg": "queryT", "semantic": "query", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" }, { "name": "k", "arg": "keyT", "semantic": "key", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" }, { "name": "v", "arg": "valueT", "semantic": "value", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" }, { "name": "y", "arg": "outputT", "semantic": "output", "buffer": { "type": "storage" }, "elementType": "$scalar" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [ { "name": "count", "type": "u32", "value": "numel(shapes.output)" }, { "name": "qHeads", "type": "u32", "value": "attrs.num_heads" }, { "name": "kvHeads", "type": "u32", "value": "attrs.num_heads" }, { "name": "qSeq", "type": "u32", "value": "dim(shapes.query, 1)" }, { "name": "kvSeq", "type": "u32", "value": "dim(shapes.key, 1)" }, { "name": "headSize", "type": "u32", "value": "dim(shapes.query, 2) / attrs.num_heads" }, { "name": "vHeadSize", "type": "u32", "value": "dim(shapes.value, 2) / attrs.num_heads" }, { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" }, { "name": "softcap", "type": "f32", "value": "0" }, { "name": "isCausal", "type": "u32", "value": "attrs.unidirectional" }, { "name": "qHidden", "type": "u32", "value": "dim(shapes.query, 2)" }, { "name": "kvHidden", "type": "u32", "value": "dim(shapes.key, 2)" }, { "name": "vHidden", "type": "u32", "value": "dim(shapes.value, 2)" } ] } } ], "prefillTiledNoBiasMask": [ { "name": "q", "arg": "queryT", "semantic": "query", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" }, { "name": "k", "arg": "keyT", "semantic": "key", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" }, { "name": "v", "arg": "valueT", "semantic": "value", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" }, { "name": "attn_mask", "arg": "attentionBiasT", "semantic": "attention_bias", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" }, { "name": "y", "arg": "outputT", "semantic": "output", "buffer": { "type": "storage" }, "elementType": "$scalar" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [ { "name": "count", "type": "u32", "value": "numel(shapes.output)" }, { "name": "qHeads", "type": "u32", "value": "attrs.num_heads" }, { "name": "kvHeads", "type": "u32", "value": "attrs.num_heads" }, { "name": "qSeq", "type": "u32", "value": "dim(shapes.query, 1)" }, { "name": "kvSeq", "type": "u32", "value": "dim(shapes.key, 1)" }, { "name": "headSize", "type": "u32", "value": "dim(shapes.query, 2) / attrs.num_heads" }, { "name": "vHeadSize", "type": "u32", "value": "dim(shapes.value, 2) / attrs.num_heads" }, { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" }, { "name": "softcap", "type": "f32", "value": "0" }, { "name": "isCausal", "type": "u32", "value": "attrs.unidirectional" }, { "name": "qHidden", "type": "u32", "value": "dim(shapes.query, 2)" }, { "name": "kvHidden", "type": "u32", "value": "dim(shapes.key, 2)" }, { "name": "vHidden", "type": "u32", "value": "dim(shapes.value, 2)" }, { "name": "maskBatchStride", "type": "u32", "value": "0 if dim(shapes.attention_bias, 0) == 1 else dim(shapes.attention_bias, 1) * dim(shapes.attention_bias, 2) * dim(shapes.attention_bias, 3)" }, { "name": "maskHeadStride", "type": "u32", "value": "0 if dim(shapes.attention_bias, 1) == 1 else dim(shapes.attention_bias, 2) * dim(shapes.attention_bias, 3)" }, { "name": "maskSeqStride", "type": "u32", "value": "dim(shapes.attention_bias, 3)" } ] } } ], "materializedApplyBias": [ { "name": "scores", "semantic": "materializedScores", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, { "name": "value", "arg": "valueT", "semantic": "value", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, { "name": "bias", "arg": "biasT", "semantic": "bias", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, { "name": "output", "arg": "outputT", "semantic": "output", "buffer": { "type": "storage" }, "elementType": "f32" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [ { "name": "qSeq", "type": "u32", "value": "dim(shapes.query, 1)" }, { "name": "kvSeq", "type": "u32", "value": "dim(shapes.key, 1)" } ] } } ], "materializedScoreStats": [ { "name": "query", "arg": "queryT", "semantic": "query", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, { "name": "key", "arg": "keyT", "semantic": "key", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, { "name": "scores", "semantic": "materializedScores", "buffer": { "type": "storage" }, "elementType": "f32" }, { "name": "scorePartials", "semantic": "materializedScorePartials", "buffer": { "type": "storage" }, "elementType": "f32" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [ { "name": "qSeq", "type": "u32", "value": "dim(shapes.query, 1)" }, { "name": "kvSeq", "type": "u32", "value": "dim(shapes.key, 1)" }, { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" } ] } } ], "materializedScoreStatsBias": [ { "name": "query", "arg": "queryT", "semantic": "query", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, { "name": "key", "arg": "keyT", "semantic": "key", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, { "name": "bias", "arg": "biasT", "semantic": "bias", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, { "name": "scores", "semantic": "materializedScores", "buffer": { "type": "storage" }, "elementType": "f32" }, { "name": "scorePartials", "semantic": "materializedScorePartials", "buffer": { "type": "storage" }, "elementType": "f32" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [ { "name": "qSeq", "type": "u32", "value": "dim(shapes.query, 1)" }, { "name": "kvSeq", "type": "u32", "value": "dim(shapes.key, 1)" }, { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" } ] } } ], "materializedRowStatsCombine": [ { "name": "scorePartials", "semantic": "materializedScorePartials", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, { "name": "rowStats", "semantic": "materializedRowStats", "buffer": { "type": "storage" }, "elementType": "f32" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [ { "name": "rows", "type": "u32", "value": "dim(shapes.query, 0) * attrs.num_heads * dim(shapes.query, 1)" } ] } } ], "materializedApplyFused": [ { "name": "scores", "semantic": "materializedScores", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, { "name": "value", "arg": "valueT", "semantic": "value", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, { "name": "output", "arg": "outputT", "semantic": "output", "buffer": { "type": "storage" }, "elementType": "f32" }, { "name": "rowStats", "semantic": "materializedRowStats", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [ { "name": "qSeq", "type": "u32", "value": "dim(shapes.query, 1)" }, { "name": "kvSeq", "type": "u32", "value": "dim(shapes.key, 1)" } ] } } ], "materializedApplyFusedBias": [ { "name": "scores", "semantic": "materializedScores", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, { "name": "value", "arg": "valueT", "semantic": "value", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, { "name": "bias", "arg": "biasT", "semantic": "bias", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, { "name": "output", "arg": "outputT", "semantic": "output", "buffer": { "type": "storage" }, "elementType": "f32" }, { "name": "rowStats", "semantic": "materializedRowStats", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [ { "name": "qSeq", "type": "u32", "value": "dim(shapes.query, 1)" }, { "name": "kvSeq", "type": "u32", "value": "dim(shapes.key, 1)" } ] } } ], "materializedScoreStatsDyn": [ { "name": "query", "arg": "queryT", "semantic": "query", "buffer": { "type": "read-only-storage" }, "elementType": "$inputElement" }, { "name": "key", "arg": "keyT", "semantic": "key", "buffer": { "type": "read-only-storage" }, "elementType": "$inputElement" }, { "name": "scores", "semantic": "materializedScores", "buffer": { "type": "storage" }, "elementType": "f32" }, { "name": "scorePartials", "semantic": "materializedScorePartials", "buffer": { "type": "storage" }, "elementType": "f32" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [ { "name": "qSeq", "type": "u32", "value": "dim(shapes.query, 1)" }, { "name": "kvSeq", "type": "u32", "value": "dim(shapes.key, 1)" }, { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" } ] } } ], "materializedApplyFusedDyn": [ { "name": "scores", "semantic": "materializedScores", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, { "name": "value", "arg": "valueT", "semantic": "value", "buffer": { "type": "read-only-storage" }, "elementType": "$inputElement" }, { "name": "output", "arg": "outputT", "semantic": "output", "buffer": { "type": "storage" }, "elementType": "$outputElement" }, { "name": "rowStats", "semantic": "materializedRowStats", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [ { "name": "qSeq", "type": "u32", "value": "dim(shapes.query, 1)" }, { "name": "kvSeq", "type": "u32", "value": "dim(shapes.key, 1)" } ] } } ] }, "variants": [ { "id": "qkv_no_bias_small_seq", "description": "Stages each head's complete key and value planes in one workgroup so short bidirectional FP32 requests do not launch work per query/head pair. It yields whenever a flash route is admissible.", "priority": 60, "when": ["not present.biasT", "smallSeqShapeOk", "not flashShapeOk", "smallSeqResourcesFit", "smallSeqDispatchFits"], "constants": { "usesF16": false, "inputElement": "\"f32\"", "outputElement": "\"f32\"", "headDim": "dim(shapes.query, 2) / attrs.num_heads", "kvSeq": "dim(shapes.key, 1)", "hidden": "dim(shapes.query, 2)", "workgroupSize": "smallSeqWorkgroupSize" }, "passes": [ { "id": "main", "name": "MultiHeadAttention", "bindings": "noBiasSmallBidirectional", "dispatch": { "x": "attrs.num_heads", "y": "dim(shapes.query, 0)" }, "profile": { "op": "\"multi_head_attention\"", "variant": "\"qkv_no_bias_small_seq\"", "numHeads": "attrs.num_heads", "headDim": "dim(shapes.query, 2) / attrs.num_heads" }, "shader": "mha-small-seq.wgsl.jinja" } ] }, { "id": "qkv_no_bias_tiled_nosg", "description": "Portable tiled prefill route for devices without an admissible subgroup path. It computes attention online without materializing the full score matrix.", "priority": 19, "supersededBy": ["qkv_no_bias_flash_cluster_nosg", "qkv_no_bias_flash_cluster_lpq4_nosg"], "when": ["not present.biasT", "qkvContractOk", "not smallHeadParallelOk", "headDim % 4 == 0", "headDim <= 128", "dim(shapes.query, 1) >= 31", "prefillTiledDeviceOk"], "constants": { "scalar": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "usesF16": "tensorDtypes.query == \"float16\"", "blockM": "tunables.PREFILL_QUERY_TILE", "vHeadCap": "dim(shapes.query, 2) / attrs.num_heads" }, "passes": [ { "id": "main", "name": "MultiHeadAttention.PrefillTiledNoSg", "source": { "shader": "attention-rank4-tiled.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, "bindings": "prefillTiledNoBias", "dispatch": { "threads": "dim(shapes.output, 0) * attrs.num_heads * ceil(dim(shapes.output, 1) / 32) * 32", "workgroupSize": "32" } } ] }, { "id": "qkv_bias_flash_q32_broadcast_f32_d128", "priority": 30, "requires": { "features": ["subgroups"] }, "when": ["tensorDtypes.query == \"float32\"", "biasOk", "attrs.unidirectional == 0", "flashShapeOk", "q32BroadcastF32RegisterGeometry", "dim(shapes.query, 1) >= 31", "ceilDiv(dim(shapes.query, 1), 32) <= device.limits.maxComputeWorkgroupsPerDimension", "subgroupsWave32"], "constants": { "hasBias": true, "hasCausal": false, "usesF16": false, "scalar": "\"f32\"", "inputVec4": "\"vec4\"", "inputElement": "\"vec4\"", "outputElement": "\"vec4\"", "inputScalar": "\"f32\"", "qNumHeads": "attrs.num_heads", "kvNumHeads": "attrs.num_heads", "qHidden": "dim(shapes.query, 2)", "headDim": "dim(shapes.query, 2) / attrs.num_heads", "headDimV4": "(dim(shapes.query, 2) / attrs.num_heads) / 4", "kStep": 64, "qkGroups": 16 }, "passes": [ { "id": "main", "name": "MultiHeadAttention.FlashQ32BroadcastF32Bias", "source": { "shader": "attn-flash-q32-broadcast.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, "bindings": "biasNonCausal", "dispatch": { "x": "ceilDiv(dim(shapes.query, 1), 32)", "y": "attrs.num_heads", "z": "dim(shapes.query, 0)" } } ] }, { "id": "qkv_no_bias_small_head_parallel", "priority": 10, "when": ["not present.biasT", "smallHeadParallelOk"], "constants": { "usesF16": "tensorDtypes.query == \"float16\"", "inputScalar": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "outputScalar": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "inputElement": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "outputElement": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "qHidden": "dim(shapes.query, 2)", "headDim": "dim(shapes.query, 2) / attrs.num_heads", "kvSeq": "dim(shapes.key, 1)" }, "passes": [ { "id": "main", "name": "MultiHeadAttentionSmallHeadParallel", "shader": "attn-small-head-parallel.wgsl.jinja", "bindings": "noBiasSmallBidirectional", "dispatch": { "x": "min(dim(shapes.query, 1), device.limits.maxComputeWorkgroupsPerDimension)", "y": "attrs.num_heads", "z": "dim(shapes.query, 0)" } } ] }, { "id": "qkv_no_bias_tiled_attn_bias_nosg", "description": "Portable tiled prefill route for additive attention bias when no subgroup path is admissible. It computes attention online without materializing the full score matrix.", "priority": 17, "when": ["not present.biasT", "qkvMaskContractOk", "not smallHeadParallelOk", "headDim % 4 == 0", "headDim <= 128", "dim(shapes.query, 1) >= 31", "prefillTiledDeviceOk"], "constants": { "scalar": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "usesF16": "tensorDtypes.query == \"float16\"", "hasMask": true, "maskIsBool": false, "blockM": "tunables.PREFILL_QUERY_TILE", "vHeadCap": "dim(shapes.query, 2) / attrs.num_heads" }, "passes": [ { "id": "main", "name": "MultiHeadAttention.PrefillTiledAttnBiasNoSg", "source": { "shader": "attention-rank4-tiled.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, "bindings": "prefillTiledNoBiasMask", "dispatch": { "threads": "dim(shapes.output, 0) * attrs.num_heads * ceil(dim(shapes.output, 1) / 32) * 32", "workgroupSize": "32" } } ] }, { "id": "qkv_no_bias_materialized_sgmat_fused_f32", "description": "Materializes FP32 scores for subgroup-matrix score and value products, and folds the softmax into the apply pass using row statistics the score pass emits. Selected over the separate-softmax route when the score matrix is large enough that reading it back costs more than the apply-side staging fusing forces.", "priority": 52, "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] }, "when": ["materializedSgmatFusedOk", "not present.biasT", "materializedFusedSoftmaxWorthIt"], "constants": { "qNumHeads": "attrs.num_heads", "headDim": "headDim", "qHidden": "dim(shapes.query, 2)", "materializedSgmatQueryTile": "materializedSgmatQueryTile", "materializedSgmatKeyTile": "materializedSgmatKeyTile", "materializedSgmatInnerTile": "materializedSgmatInnerTile", "hasBias": false, "useSubgroups": true, "materializedRowStatsWg": "materializedRowStatsWg", "statSlots": "materializedSgmatStatSlots", "statQuerySeq": "dim(shapes.query, 1)" }, "intermediates": [ { "id": "materializedScores", "dtype": "float32", "shape": "[dim(shapes.query, 0) * attrs.num_heads * dim(shapes.query, 1) * dim(shapes.key, 1)]" }, { "id": "materializedRowStats", "dtype": "float32", "shape": "[materializedRowStatsElements]" }, { "id": "materializedScorePartials", "dtype": "float32", "shape": "[materializedScorePartialElements]" } ], "passes": [ { "id": "scores", "name": "MultiHeadAttention.MaterializedScoresSgmat", "source": { "shader": "attn-materialized-sgmat-f32.wgsl.jinja", "inputs": { "phase": "\"score\"", "emitRowStats": true } }, "bindings": "materializedScoreStats", "dispatch": { "x": "ceilDiv(dim(shapes.key, 1), materializedSgmatKeyTile)", "y": "ceilDiv(dim(shapes.query, 1), materializedSgmatQueryTile)", "z": "dim(shapes.query, 0) * attrs.num_heads" } }, { "id": "rowstats", "name": "MultiHeadAttention.MaterializedRowStatsCombine", "source": { "shader": "attn-materialized-rowstats-combine-f32.wgsl.jinja" }, "bindings": "materializedRowStatsCombine", "dispatch": { "gridStride": "dim(shapes.query, 0) * attrs.num_heads * dim(shapes.query, 1)", "workgroupSize": "materializedRowStatsWg" } }, { "id": "apply", "name": "MultiHeadAttention.MaterializedApplySgmat", "source": { "shader": "attn-materialized-sgmat-f32.wgsl.jinja", "inputs": { "phase": "\"apply\"", "fusedSoftmax": true } }, "bindings": "materializedApplyFused", "dispatch": { "x": "ceilDiv(headDim, materializedSgmatKeyTile)", "y": "ceilDiv(dim(shapes.query, 1), materializedSgmatQueryTile)", "z": "dim(shapes.query, 0) * attrs.num_heads" } } ] }, { "id": "qkv_bias_materialized_sgmat_fused_f32", "description": "Materializes FP32 scores for subgroup-matrix score and value products, and folds the softmax into the apply pass using row statistics the score pass emits. Selected over the separate-softmax route when the score matrix is large enough that reading it back costs more than the apply-side staging fusing forces.", "priority": 52, "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] }, "when": ["materializedSgmatFusedOk", "biasOk", "materializedFusedSoftmaxWorthIt"], "constants": { "qNumHeads": "attrs.num_heads", "headDim": "headDim", "qHidden": "dim(shapes.query, 2)", "materializedSgmatQueryTile": "materializedSgmatQueryTile", "materializedSgmatKeyTile": "materializedSgmatKeyTile", "materializedSgmatInnerTile": "materializedSgmatInnerTile", "hasBias": true, "useSubgroups": true, "materializedRowStatsWg": "materializedRowStatsWg", "statSlots": "materializedSgmatStatSlots", "statQuerySeq": "dim(shapes.query, 1)" }, "intermediates": [ { "id": "materializedScores", "dtype": "float32", "shape": "[dim(shapes.query, 0) * attrs.num_heads * dim(shapes.query, 1) * dim(shapes.key, 1)]" }, { "id": "materializedRowStats", "dtype": "float32", "shape": "[materializedRowStatsElements]" }, { "id": "materializedScorePartials", "dtype": "float32", "shape": "[materializedScorePartialElements]" } ], "passes": [ { "id": "scores", "name": "MultiHeadAttention.MaterializedScoresSgmatBias", "source": { "shader": "attn-materialized-sgmat-f32.wgsl.jinja", "inputs": { "phase": "\"score\"", "emitRowStats": true } }, "bindings": "materializedScoreStatsBias", "dispatch": { "x": "ceilDiv(dim(shapes.key, 1), materializedSgmatKeyTile)", "y": "ceilDiv(dim(shapes.query, 1), materializedSgmatQueryTile)", "z": "dim(shapes.query, 0) * attrs.num_heads" } }, { "id": "rowstats", "name": "MultiHeadAttention.MaterializedRowStatsCombineBias", "source": { "shader": "attn-materialized-rowstats-combine-f32.wgsl.jinja" }, "bindings": "materializedRowStatsCombine", "dispatch": { "gridStride": "dim(shapes.query, 0) * attrs.num_heads * dim(shapes.query, 1)", "workgroupSize": "materializedRowStatsWg" } }, { "id": "apply", "name": "MultiHeadAttention.MaterializedApplySgmatBias", "source": { "shader": "attn-materialized-sgmat-f32.wgsl.jinja", "inputs": { "phase": "\"apply\"", "fusedSoftmax": true } }, "bindings": "materializedApplyFusedBias", "dispatch": { "x": "ceilDiv(headDim, materializedSgmatKeyTile)", "y": "ceilDiv(dim(shapes.query, 1), materializedSgmatQueryTile)", "z": "dim(shapes.query, 0) * attrs.num_heads" } } ] }, { "id": "qkv_no_bias_materialized_sgmat_fused_f16", "description": "Materializes FP32 scores for subgroup-matrix score and value products, and folds the softmax into the apply pass using row statistics the score pass emits. Selected over the separate-softmax route when the score matrix is large enough that reading it back costs more than the apply-side staging fusing forces.", "priority": 52, "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }] }, "when": ["materializedSgmatFusedF16Ok", "not present.biasT", "materializedFusedSoftmaxWorthIt"], "constants": { "qNumHeads": "attrs.num_heads", "headDim": "headDim", "qHidden": "dim(shapes.query, 2)", "materializedSgmatQueryTile": "materializedSgmatQueryTile", "materializedSgmatKeyTile": "materializedSgmatKeyTile", "materializedSgmatInnerTile": "materializedSgmatInnerTile", "hasBias": false, "useSubgroups": true, "materializedRowStatsWg": "materializedRowStatsWg", "statSlots": "materializedSgmatStatSlots", "statQuerySeq": "dim(shapes.query, 1)", "operandF16": true, "inputElement": "\"f16\"", "outputElement": "\"f16\"" }, "intermediates": [ { "id": "materializedScores", "dtype": "float32", "shape": "[dim(shapes.query, 0) * attrs.num_heads * dim(shapes.query, 1) * dim(shapes.key, 1)]" }, { "id": "materializedRowStats", "dtype": "float32", "shape": "[materializedRowStatsElements]" }, { "id": "materializedScorePartials", "dtype": "float32", "shape": "[materializedScorePartialElements]" } ], "passes": [ { "id": "scores", "name": "MultiHeadAttention.MaterializedScoresSgmatF16", "source": { "shader": "attn-materialized-sgmat-f32.wgsl.jinja", "inputs": { "phase": "\"score\"", "emitRowStats": true } }, "bindings": "materializedScoreStatsDyn", "dispatch": { "x": "ceilDiv(dim(shapes.key, 1), materializedSgmatKeyTile)", "y": "ceilDiv(dim(shapes.query, 1), materializedSgmatQueryTile)", "z": "dim(shapes.query, 0) * attrs.num_heads" } }, { "id": "rowstats", "name": "MultiHeadAttention.MaterializedRowStatsCombine", "source": { "shader": "attn-materialized-rowstats-combine-f32.wgsl.jinja" }, "bindings": "materializedRowStatsCombine", "dispatch": { "gridStride": "dim(shapes.query, 0) * attrs.num_heads * dim(shapes.query, 1)", "workgroupSize": "materializedRowStatsWg" } }, { "id": "apply", "name": "MultiHeadAttention.MaterializedApplySgmatF16", "source": { "shader": "attn-materialized-sgmat-f32.wgsl.jinja", "inputs": { "phase": "\"apply\"", "fusedSoftmax": true } }, "bindings": "materializedApplyFusedDyn", "dispatch": { "x": "ceilDiv(headDim, materializedSgmatKeyTile)", "y": "ceilDiv(dim(shapes.query, 1), materializedSgmatQueryTile)", "z": "dim(shapes.query, 0) * attrs.num_heads" } } ] }, { "id": "qkv_no_bias_flash_cluster_lpq4_nosg", "description": "Keeps query/key tiles and online-softmax state inside one workgroup instead of materializing all scores. The family supplies subgroup and portable forms for compatible prefill shapes.", "priority": 20, "requires": {}, "when": ["not present.biasT", "flashShapeOk", "headDim % 16 == 0", "headDim % 32 != 0", "dim(shapes.query, 1) >= 31"], "constants": { "hasBias": false, "hasCausal": true, "usesF16": "tensorDtypes.query == \"float16\"", "scalar": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "inputVec4": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "inputElement": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "outputElement": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "qNumHeads": "attrs.num_heads", "kvNumHeads": "attrs.num_heads", "qHiddenV4": "dim(shapes.query, 2) / 4", "kvHiddenV4": "dim(shapes.query, 2) / 4", "headDim": "dim(shapes.query, 2) / attrs.num_heads", "headDimV4": "(dim(shapes.query, 2) / attrs.num_heads) / 4", "TILE_Q": 16, "TILE_K": "clusterTileKWg64", "batchNoSgReduction": "tensorDtypes.query == \"float16\"", "LPQ": 4, "useSubgroups": false }, "passes": [ { "id": "main", "name": "MultiHeadAttention.Flash", "source": { "shader": "attn-flash-prefill-cluster.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, "bindings": "noBias", "dispatch": { "x": "ceilDiv(dim(shapes.query, 1), constants.TILE_Q)", "y": "attrs.num_heads", "z": "dim(shapes.query, 0)" } } ] }, { "id": "qkv_no_bias_flash_cluster_nosg", "description": "Keeps query/key tiles and online-softmax state inside one workgroup instead of materializing all scores. The family supplies subgroup and portable forms for compatible prefill shapes.", "priority": 20, "requires": {}, "when": ["not present.biasT", "flashShapeOk", "headDim % 32 == 0", "dim(shapes.query, 1) >= 31"], "constants": { "hasBias": false, "hasCausal": true, "usesF16": "tensorDtypes.query == \"float16\"", "scalar": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "inputVec4": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "inputElement": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "outputElement": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "qNumHeads": "attrs.num_heads", "kvNumHeads": "attrs.num_heads", "qHiddenV4": "dim(shapes.query, 2) / 4", "kvHiddenV4": "dim(shapes.query, 2) / 4", "headDim": "dim(shapes.query, 2) / attrs.num_heads", "headDimV4": "(dim(shapes.query, 2) / attrs.num_heads) / 4", "TILE_Q": 16, "TILE_K": "clusterTileKWg128", "batchNoSgReduction": "tensorDtypes.query == \"float16\"", "LPQ": 8, "useSubgroups": false }, "passes": [ { "id": "main", "name": "MultiHeadAttention.Flash", "source": { "shader": "attn-flash-prefill-cluster.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, "bindings": "noBias", "dispatch": { "x": "ceilDiv(dim(shapes.query, 1), constants.TILE_Q)", "y": "attrs.num_heads", "z": "dim(shapes.query, 0)" } } ] }, { "id": "qkv_bias_flash_cluster_nosg", "description": "Keeps query/key tiles and online-softmax state inside one workgroup instead of materializing all scores. The family supplies subgroup and portable forms for compatible prefill shapes.", "priority": 19, "requires": {}, "when": ["biasOk", "flashShapeOk", "headDim % 32 == 0", "dim(shapes.query, 1) >= 31"], "constants": { "hasBias": true, "hasCausal": true, "usesF16": "tensorDtypes.query == \"float16\"", "scalar": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "inputVec4": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "inputElement": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "outputElement": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "qNumHeads": "attrs.num_heads", "kvNumHeads": "attrs.num_heads", "qHiddenV4": "dim(shapes.query, 2) / 4", "kvHiddenV4": "dim(shapes.query, 2) / 4", "headDim": "dim(shapes.query, 2) / attrs.num_heads", "headDimV4": "(dim(shapes.query, 2) / attrs.num_heads) / 4", "TILE_Q": 16, "TILE_K": "clusterTileKWg128", "batchNoSgReduction": "tensorDtypes.query == \"float16\"", "qHidden": "dim(shapes.query, 2)", "inputScalar": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "LPQ": 8, "useSubgroups": false }, "passes": [ { "id": "main", "name": "MultiHeadAttention.Flash", "source": { "shader": "attn-flash-prefill-cluster.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, "bindings": "bias", "dispatch": { "x": "ceilDiv(dim(shapes.query, 1), constants.TILE_Q)", "y": "attrs.num_heads", "z": "dim(shapes.query, 0)" } } ] }, { "id": "qkv_no_bias_flash_cluster_lpq4", "description": "Keeps query/key tiles and online-softmax state inside one workgroup instead of materializing all scores. The family supplies subgroup and portable forms for compatible prefill shapes.", "priority": 22, "requires": { "features": ["subgroups"] }, "when": ["not present.biasT", "flashShapeOk", "headDim % 16 == 0", "headDim % 32 != 0", "dim(shapes.query, 1) >= 31", "subgroupCluster4"], "constants": { "hasBias": false, "hasCausal": true, "usesF16": "tensorDtypes.query == \"float16\"", "scalar": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "inputVec4": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "inputElement": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "outputElement": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "qNumHeads": "attrs.num_heads", "kvNumHeads": "attrs.num_heads", "qHiddenV4": "dim(shapes.query, 2) / 4", "kvHiddenV4": "dim(shapes.query, 2) / 4", "headDim": "dim(shapes.query, 2) / attrs.num_heads", "headDimV4": "(dim(shapes.query, 2) / attrs.num_heads) / 4", "TILE_Q": 16, "TILE_K": "32 if (tensorDtypes.query != \"float16\" and (dim(shapes.query, 2) / attrs.num_heads) <= 64) else 8", "LPQ": 4 }, "passes": [ { "id": "main", "name": "MultiHeadAttention.Flash", "source": { "shader": "attn-flash-prefill-cluster.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, "bindings": "noBias", "dispatch": { "x": "ceilDiv(dim(shapes.query, 1), constants.TILE_Q)", "y": "attrs.num_heads", "z": "dim(shapes.query, 0)" } } ] }, { "id": "qkv_no_bias_flash_cluster", "description": "Keeps query/key tiles and online-softmax state inside one workgroup instead of materializing all scores. The family supplies subgroup and portable forms for compatible prefill shapes.", "priority": 22, "requires": { "features": ["subgroups"] }, "when": ["not present.biasT", "flashShapeOk", "headDim % 32 == 0", "dim(shapes.query, 1) >= 31", "subgroupCluster8"], "constants": { "hasBias": false, "hasCausal": true, "usesF16": "tensorDtypes.query == \"float16\"", "scalar": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "inputVec4": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "inputElement": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "outputElement": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "qNumHeads": "attrs.num_heads", "kvNumHeads": "attrs.num_heads", "qHiddenV4": "dim(shapes.query, 2) / 4", "kvHiddenV4": "dim(shapes.query, 2) / 4", "headDim": "dim(shapes.query, 2) / attrs.num_heads", "headDimV4": "(dim(shapes.query, 2) / attrs.num_heads) / 4", "TILE_Q": 16, "TILE_K": "32 if (tensorDtypes.query != \"float16\" and (dim(shapes.query, 2) / attrs.num_heads) <= 64) else 8", "LPQ": 8 }, "passes": [ { "id": "main", "name": "MultiHeadAttention.Flash", "source": { "shader": "attn-flash-prefill-cluster.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, "bindings": "noBias", "dispatch": { "x": "ceilDiv(dim(shapes.query, 1), constants.TILE_Q)", "y": "attrs.num_heads", "z": "dim(shapes.query, 0)" } } ] }, { "id": "qkv_bias_flash_cluster", "description": "Keeps query/key tiles and online-softmax state inside one workgroup instead of materializing all scores. The family supplies subgroup and portable forms for compatible prefill shapes.", "priority": 21, "requires": { "features": ["subgroups"] }, "when": ["biasOk", "flashShapeOk", "headDim % 32 == 0", "dim(shapes.query, 1) >= 31", "subgroupCluster8"], "constants": { "hasBias": true, "hasCausal": true, "usesF16": "tensorDtypes.query == \"float16\"", "scalar": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "inputVec4": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "inputElement": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "outputElement": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "qNumHeads": "attrs.num_heads", "kvNumHeads": "attrs.num_heads", "qHiddenV4": "dim(shapes.query, 2) / 4", "kvHiddenV4": "dim(shapes.query, 2) / 4", "headDim": "dim(shapes.query, 2) / attrs.num_heads", "headDimV4": "(dim(shapes.query, 2) / attrs.num_heads) / 4", "TILE_Q": 16, "TILE_K": "32 if (tensorDtypes.query != \"float16\" and (dim(shapes.query, 2) / attrs.num_heads) <= 64) else 8", "qHidden": "dim(shapes.query, 2)", "inputScalar": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "LPQ": 8 }, "passes": [ { "id": "main", "name": "MultiHeadAttention.Flash", "source": { "shader": "attn-flash-prefill-cluster.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, "bindings": "bias", "dispatch": { "x": "ceilDiv(dim(shapes.query, 1), constants.TILE_Q)", "y": "attrs.num_heads", "z": "dim(shapes.query, 0)" } } ] }, { "id": "qkv_no_bias_flash_cluster_attn_bias", "description": "Keeps query/key tiles and online-softmax state inside one workgroup instead of materializing all scores. The family supplies subgroup and portable forms for compatible prefill shapes.", "priority": 22, "requires": { "features": ["subgroups"] }, "when": ["not present.biasT", "flashMaskShapeOk", "headDim % 32 == 0", "dim(shapes.query, 1) >= 31", "subgroupCluster8"], "constants": { "hasBias": false, "hasCausal": true, "hasMask": true, "maskIsBool": false, "usesF16": "tensorDtypes.query == \"float16\"", "scalar": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "inputVec4": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "inputElement": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "outputElement": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "maskElement": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "qNumHeads": "attrs.num_heads", "kvNumHeads": "attrs.num_heads", "qHiddenV4": "dim(shapes.query, 2) / 4", "kvHiddenV4": "dim(shapes.query, 2) / 4", "headDim": "dim(shapes.query, 2) / attrs.num_heads", "headDimV4": "(dim(shapes.query, 2) / attrs.num_heads) / 4", "TILE_Q": 16, "TILE_K": "32 if (tensorDtypes.query != \"float16\" and (dim(shapes.query, 2) / attrs.num_heads) <= 64) else 8", "LPQ": 8 }, "passes": [ { "id": "main", "name": "MultiHeadAttention.Flash", "source": { "shader": "attn-flash-prefill-cluster.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, "bindings": "noBiasMask", "dispatch": { "x": "ceilDiv(dim(shapes.query, 1), constants.TILE_Q)", "y": "attrs.num_heads", "z": "dim(shapes.query, 0)" } } ] }, { "id": "qkv_bias_flash_cluster_attn_bias", "description": "Keeps query/key tiles and online-softmax state inside one workgroup instead of materializing all scores. The family supplies subgroup and portable forms for compatible prefill shapes.", "priority": 21, "requires": { "features": ["subgroups"] }, "when": ["biasOk", "flashMaskShapeOk", "headDim % 32 == 0", "dim(shapes.query, 1) >= 31", "subgroupCluster8"], "constants": { "hasBias": true, "hasCausal": true, "hasMask": true, "maskIsBool": false, "usesF16": "tensorDtypes.query == \"float16\"", "scalar": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "inputVec4": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "inputElement": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "outputElement": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "maskElement": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "qNumHeads": "attrs.num_heads", "kvNumHeads": "attrs.num_heads", "qHiddenV4": "dim(shapes.query, 2) / 4", "kvHiddenV4": "dim(shapes.query, 2) / 4", "headDim": "dim(shapes.query, 2) / attrs.num_heads", "headDimV4": "(dim(shapes.query, 2) / attrs.num_heads) / 4", "TILE_Q": 16, "TILE_K": "32 if (tensorDtypes.query != \"float16\" and (dim(shapes.query, 2) / attrs.num_heads) <= 64) else 8", "qHidden": "dim(shapes.query, 2)", "inputScalar": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "LPQ": 8 }, "passes": [ { "id": "main", "name": "MultiHeadAttention.Flash", "source": { "shader": "attn-flash-prefill-cluster.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, "bindings": "biasMask", "dispatch": { "x": "ceilDiv(dim(shapes.query, 1), constants.TILE_Q)", "y": "attrs.num_heads", "z": "dim(shapes.query, 0)" } } ] }, { "id": "qkv_no_bias_flash_splitk_nosg", "description": "Partitions the key reduction across workgroups and merges partial softmax statistics and values. Used when decode or another short query has too few independent query/head workgroups; the merge reassociates the key reduction.", "priority": 21, "requires": {}, "when": ["not present.biasT", "decodeSplitKNoBiasOk or shortQuerySplitKNoBiasOk"], "constants": { "combineSubgroups": false, "useSubgroups": false, "splitQueries": true, "scalar": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "inputVec4": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "qNumHeads": "attrs.num_heads", "kvNumHeads": "attrs.num_heads", "headDim": "dim(shapes.query, 2) / attrs.num_heads", "headDimV4": "(dim(shapes.query, 2) / attrs.num_heads) / 4", "qSeq": "dim(shapes.query, 1)", "qHiddenV4": "dim(shapes.query, 2) / 4", "kvHiddenV4": "dim(shapes.query, 2) / 4", "numSplits": "noBiasSplitKCount", "usesF16": "tensorDtypes.query == \"float16\"" }, "intermediates": [ { "id": "partialOut", "dtype": "float32", "shape": "[dim(shapes.query, 0) * dim(shapes.query, 1) * attrs.num_heads * noBiasSplitKCount * (dim(shapes.query, 2) / attrs.num_heads)]" }, { "id": "partialStats", "dtype": "float32", "shape": "[2 * dim(shapes.query, 0) * dim(shapes.query, 1) * attrs.num_heads * noBiasSplitKCount]" } ], "passes": [ { "id": "split_attention", "name": "MultiHeadAttention.DecodeSplitKNoSg", "source": { "shader": "attn-flash-decode-splitk.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, "bindings": "decodeSplitPartial", "dispatch": { "x": "dim(shapes.query, 1) * noBiasSplitKCount", "y": "attrs.num_heads", "z": "dim(shapes.query, 0)" } }, { "id": "merge", "name": "MultiHeadAttention.DecodeSplitKMergeNoSg", "source": { "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, "bindings": "decodeSplitMergeNoParams", "dispatch": { "x": "dim(shapes.query, 1)", "y": "attrs.num_heads", "z": "dim(shapes.query, 0)" } } ] }, { "id": "qkv_bias_flash_splitk_nosg", "description": "Partitions the key reduction across workgroups and merges partial softmax statistics and values. Used when decode or another short query has too few independent query/head workgroups; the merge reassociates the key reduction.", "priority": 17, "requires": {}, "when": ["biasOk", "decodeSplitKBiasOk"], "constants": { "combineSubgroups": false, "useSubgroups": false, "hasBias": true, "scalar": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "inputVec4": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "inputScalar": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "qNumHeads": "attrs.num_heads", "kvNumHeads": "attrs.num_heads", "qHidden": "dim(shapes.query, 2)", "headDim": "dim(shapes.query, 2) / attrs.num_heads", "headDimV4": "(dim(shapes.query, 2) / attrs.num_heads) / 4", "qHiddenV4": "dim(shapes.query, 2) / 4", "kvHiddenV4": "dim(shapes.query, 2) / 4", "numSplits": "biasSplitKCount", "usesF16": "tensorDtypes.query == \"float16\"" }, "intermediates": [ { "id": "partialOut", "dtype": "float32", "shape": "[dim(shapes.query, 0) * attrs.num_heads * biasSplitKCount * (dim(shapes.query, 2) / attrs.num_heads)]" }, { "id": "partialStats", "dtype": "float32", "shape": "[2 * dim(shapes.query, 0) * attrs.num_heads * biasSplitKCount]" } ], "passes": [ { "id": "split_attention", "name": "MultiHeadAttention.DecodeSplitKBiasNoSg", "source": { "shader": "attn-flash-decode-splitk.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, "bindings": "decodeSplitBiasPartial", "dispatch": { "x": "biasSplitKCount", "y": "attrs.num_heads", "z": "dim(shapes.query, 0)" } }, { "id": "merge", "name": "MultiHeadAttention.DecodeSplitKMergeBiasNoSg", "source": { "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, "bindings": "decodeSplitBiasMergeNoParams", "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.query, 0)" } } ] }, { "id": "qkv_no_bias_flash_splitk", "description": "Partitions the key reduction across workgroups and merges partial softmax statistics and values. Used when decode or another short query has too few independent query/head workgroups; the merge reassociates the key reduction.", "priority": 25, "requires": { "features": ["subgroups"] }, "demoteWhen": ["decodeSplitKPortablePreferred"], "when": ["not present.biasT", "decodeSplitKNoBiasOk"], "constants": { "combineSubgroups": true, "splitQueries": true, "scalar": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "inputVec4": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "qNumHeads": "attrs.num_heads", "kvNumHeads": "attrs.num_heads", "headDim": "dim(shapes.query, 2) / attrs.num_heads", "headDimV4": "(dim(shapes.query, 2) / attrs.num_heads) / 4", "qSeq": "dim(shapes.query, 1)", "qHiddenV4": "dim(shapes.query, 2) / 4", "kvHiddenV4": "dim(shapes.query, 2) / 4", "numSplits": "noBiasSplitKCount", "usesF16": "tensorDtypes.query == \"float16\"" }, "intermediates": [ { "id": "partialOut", "dtype": "float32", "shape": "[dim(shapes.query, 0) * dim(shapes.query, 1) * attrs.num_heads * noBiasSplitKCount * (dim(shapes.query, 2) / attrs.num_heads)]" }, { "id": "partialStats", "dtype": "float32", "shape": "[2 * dim(shapes.query, 0) * dim(shapes.query, 1) * attrs.num_heads * noBiasSplitKCount]" } ], "passes": [ { "id": "split_attention", "name": "MultiHeadAttention.DecodeSplitK", "source": { "shader": "attn-flash-decode-splitk.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, "bindings": "decodeSplitPartial", "dispatch": { "x": "dim(shapes.query, 1) * noBiasSplitKCount", "y": "attrs.num_heads", "z": "dim(shapes.query, 0)" } }, { "id": "merge", "name": "MultiHeadAttention.DecodeSplitKMerge", "source": { "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, "bindings": "decodeSplitMergeNoParams", "dispatch": { "x": "dim(shapes.query, 1)", "y": "attrs.num_heads", "z": "dim(shapes.query, 0)" } } ] }, { "id": "qkv_bias_flash_splitk", "description": "Partitions the key reduction across workgroups and merges partial softmax statistics and values. Used when decode or another short query has too few independent query/head workgroups; the merge reassociates the key reduction.", "priority": 24, "requires": { "features": ["subgroups"] }, "when": ["biasOk", "decodeSplitKBiasOk"], "constants": { "combineSubgroups": true, "hasBias": true, "scalar": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "inputVec4": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "inputScalar": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "qNumHeads": "attrs.num_heads", "kvNumHeads": "attrs.num_heads", "qHidden": "dim(shapes.query, 2)", "headDim": "dim(shapes.query, 2) / attrs.num_heads", "headDimV4": "(dim(shapes.query, 2) / attrs.num_heads) / 4", "qHiddenV4": "dim(shapes.query, 2) / 4", "kvHiddenV4": "dim(shapes.query, 2) / 4", "numSplits": "biasSplitKCount", "usesF16": "tensorDtypes.query == \"float16\"" }, "intermediates": [ { "id": "partialOut", "dtype": "float32", "shape": "[dim(shapes.query, 0) * attrs.num_heads * biasSplitKCount * (dim(shapes.query, 2) / attrs.num_heads)]" }, { "id": "partialStats", "dtype": "float32", "shape": "[2 * dim(shapes.query, 0) * attrs.num_heads * biasSplitKCount]" } ], "passes": [ { "id": "split_attention", "name": "MultiHeadAttention.DecodeSplitKBias", "source": { "shader": "attn-flash-decode-splitk.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, "bindings": "decodeSplitBiasPartial", "dispatch": { "x": "biasSplitKCount", "y": "attrs.num_heads", "z": "dim(shapes.query, 0)" } }, { "id": "merge", "name": "MultiHeadAttention.DecodeSplitKMergeBias", "source": { "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, "bindings": "decodeSplitBiasMergeNoParams", "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.query, 0)" } } ] }, { "id": "qkv_no_bias_materialized_gemm_f32", "description": "Materializes FP32 scores and applies tiled score, softmax, and value passes without subgroup-matrix support. It covers shapes where the materialized algorithm is suitable but its capability-gated sibling is unavailable.", "priority": 42, "demoteWhen": ["narrowSubgroupRange"], "when": ["materializedF32Ok", "materializedWideSimdOk", "not present.biasT", "not materializedGemmFusedSoftmaxWorthIt"], "constants": { "qNumHeads": "attrs.num_heads", "kvNumHeads": "attrs.num_heads", "headDim": "headDim", "materializedInnerTile": "tunables.MATERIALIZED_INNER_TILE", "materializedQueryTile": "tunables.MATERIALIZED_QUERY_TILE", "materializedKeyTile": "tunables.MATERIALIZED_KEY_TILE", "materializedWorkgroupDim": "tunables.MATERIALIZED_WORKGROUP_DIM", "materializedSoftmaxWg": "materializedCachedSoftmaxWg if materializedCachedSoftmaxOk else tunables.MATERIALIZED_SOFTMAX_WORKGROUP_SIZE", "materializedSoftmaxCols": "dim(shapes.key, 1)", "materializedSoftmaxCols4": "dim(shapes.key, 1) / 4", "softmaxElementType": "\"vec4\" if materializedCachedSoftmaxOk else \"f32\"", "applyTileN": "materializedApplyTileN", "hasBias": false, "useSubgroups": "device.features.has(\"subgroups\")" }, "intermediates": [ { "id": "materializedScores", "dtype": "float32", "shape": "[dim(shapes.query, 0) * attrs.num_heads * dim(shapes.query, 1) * dim(shapes.key, 1)]" } ], "passes": [ { "id": "scores", "name": "MultiHeadAttention.MaterializedScores", "source": { "shader": "attn-materialized-score-f32.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, "bindings": "materializedScore", "dispatch": { "x": "ceilDiv(dim(shapes.key, 1), tunables.MATERIALIZED_KEY_TILE)", "y": "ceilDiv(dim(shapes.query, 1), tunables.MATERIALIZED_QUERY_TILE)", "z": "dim(shapes.query, 0) * attrs.num_heads" } }, { "id": "softmax", "name": "MultiHeadAttention.MaterializedSoftmax", "source": { "shader": "attn-materialized-softmax-f32.wgsl.jinja", "inputs": { "cacheVec4": "materializedCachedSoftmaxOk" } }, "bindings": "materializedSoftmax", "dispatch": { "workgroups": "dim(shapes.query, 0) * attrs.num_heads * dim(shapes.query, 1)" } }, { "id": "apply", "name": "MultiHeadAttention.MaterializedApply", "source": { "shader": "attn-materialized-apply-f32.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, "bindings": "materializedApply", "dispatch": { "x": "ceilDiv(headDim, constants.applyTileN)", "y": "ceilDiv(dim(shapes.query, 1), tunables.MATERIALIZED_QUERY_TILE)", "z": "dim(shapes.query, 0) * attrs.num_heads" } } ] }, { "id": "qkv_bias_materialized_gemm_f32", "description": "Materializes FP32 scores and applies tiled score, softmax, and value passes without subgroup-matrix support. It covers shapes where the materialized algorithm is suitable but its capability-gated sibling is unavailable.", "priority": 42, "when": ["materializedF32Ok", "materializedWideSimdOk", "biasOk", "not materializedGemmFusedSoftmaxWorthIt"], "constants": { "qNumHeads": "attrs.num_heads", "kvNumHeads": "attrs.num_heads", "headDim": "headDim", "materializedInnerTile": "tunables.MATERIALIZED_INNER_TILE", "materializedQueryTile": "tunables.MATERIALIZED_QUERY_TILE", "materializedKeyTile": "tunables.MATERIALIZED_KEY_TILE", "materializedWorkgroupDim": "tunables.MATERIALIZED_WORKGROUP_DIM", "materializedSoftmaxWg": "materializedCachedSoftmaxWg if materializedCachedSoftmaxOk else tunables.MATERIALIZED_SOFTMAX_WORKGROUP_SIZE", "materializedSoftmaxCols": "dim(shapes.key, 1)", "materializedSoftmaxCols4": "dim(shapes.key, 1) / 4", "softmaxElementType": "\"vec4\" if materializedCachedSoftmaxOk else \"f32\"", "applyTileN": "materializedApplyTileN", "hasBias": true, "useSubgroups": "device.features.has(\"subgroups\")" }, "intermediates": [ { "id": "materializedScores", "dtype": "float32", "shape": "[dim(shapes.query, 0) * attrs.num_heads * dim(shapes.query, 1) * dim(shapes.key, 1)]" } ], "passes": [ { "id": "scores", "name": "MultiHeadAttention.MaterializedScoresBias", "source": { "shader": "attn-materialized-score-f32.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, "bindings": "materializedScoreBias", "dispatch": { "x": "ceilDiv(dim(shapes.key, 1), tunables.MATERIALIZED_KEY_TILE)", "y": "ceilDiv(dim(shapes.query, 1), tunables.MATERIALIZED_QUERY_TILE)", "z": "dim(shapes.query, 0) * attrs.num_heads" } }, { "id": "softmax", "name": "MultiHeadAttention.MaterializedSoftmaxBias", "source": { "shader": "attn-materialized-softmax-f32.wgsl.jinja", "inputs": { "cacheVec4": "materializedCachedSoftmaxOk" } }, "bindings": "materializedSoftmax", "dispatch": { "workgroups": "dim(shapes.query, 0) * attrs.num_heads * dim(shapes.query, 1)" } }, { "id": "apply", "name": "MultiHeadAttention.MaterializedApplyBias", "source": { "shader": "attn-materialized-apply-f32.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, "bindings": "materializedApplyBias", "dispatch": { "x": "ceilDiv(headDim, constants.applyTileN)", "y": "ceilDiv(dim(shapes.query, 1), tunables.MATERIALIZED_QUERY_TILE)", "z": "dim(shapes.query, 0) * attrs.num_heads" } } ] }, { "id": "qkv_no_bias_materialized_gemm_fused_f32", "description": "Materializes FP32 scores while publishing each row's maximum, combines those row statistics, then normalizes and applies the scores in one pass. It replaces the separate softmax stage's full read and rewrite of the score matrix, which is worth its extra row-statistics pass only once that matrix is large.", "priority": 43, "demoteWhen": ["narrowSubgroupRange"], "when": ["materializedF32Ok", "materializedWideSimdOk", "not present.biasT", "materializedGemmFusedSoftmaxWorthIt"], "constants": { "qNumHeads": "attrs.num_heads", "kvNumHeads": "attrs.num_heads", "headDim": "headDim", "materializedInnerTile": "tunables.MATERIALIZED_INNER_TILE", "materializedQueryTile": "tunables.MATERIALIZED_QUERY_TILE", "materializedKeyTile": "tunables.MATERIALIZED_KEY_TILE", "materializedWorkgroupDim": "tunables.MATERIALIZED_WORKGROUP_DIM", "applyTileN": "materializedApplyTileN", "hasBias": false, "useSubgroups": "device.features.has(\"subgroups\")", "statSlots": "materializedGemmStatSlots", "statQuerySeq": "dim(shapes.query, 1)", "materializedRowStatsWg": "materializedRowStatsWg" }, "intermediates": [ { "id": "materializedScores", "dtype": "float32", "shape": "[dim(shapes.query, 0) * attrs.num_heads * dim(shapes.query, 1) * dim(shapes.key, 1)]" }, { "id": "materializedRowStats", "dtype": "float32", "shape": "[materializedRowStatsElements]" }, { "id": "materializedScorePartials", "dtype": "float32", "shape": "[materializedGemmScorePartialElements]" } ], "passes": [ { "id": "scores", "name": "MultiHeadAttention.MaterializedScoresFused", "source": { "shader": "attn-materialized-score-f32.wgsl.jinja", "inputs": { "layout": "\"bsh\"", "emitRowStats": true } }, "bindings": "materializedScoreStats", "dispatch": { "x": "ceilDiv(dim(shapes.key, 1), tunables.MATERIALIZED_KEY_TILE)", "y": "ceilDiv(dim(shapes.query, 1), tunables.MATERIALIZED_QUERY_TILE)", "z": "dim(shapes.query, 0) * attrs.num_heads" } }, { "id": "rowstats", "name": "MultiHeadAttention.MaterializedGemmRowStatsCombine", "source": { "shader": "attn-materialized-rowstats-combine-f32.wgsl.jinja", "inputs": { "maxOnly": true } }, "bindings": "materializedRowStatsCombine", "dispatch": { "gridStride": "dim(shapes.query, 0) * attrs.num_heads * dim(shapes.query, 1)", "workgroupSize": "materializedRowStatsWg" } }, { "id": "apply", "name": "MultiHeadAttention.MaterializedApplyFused", "source": { "shader": "attn-materialized-apply-f32.wgsl.jinja", "inputs": { "layout": "\"bsh\"", "fusedSoftmax": true } }, "bindings": "materializedApplyFused", "dispatch": { "x": "ceilDiv(headDim, constants.applyTileN)", "y": "ceilDiv(dim(shapes.query, 1), tunables.MATERIALIZED_QUERY_TILE)", "z": "dim(shapes.query, 0) * attrs.num_heads" } } ] }, { "id": "qkv_bias_materialized_gemm_fused_f32", "description": "Materializes FP32 scores while publishing each row's maximum, combines those row statistics, then normalizes and applies the scores in one pass. It replaces the separate softmax stage's full read and rewrite of the score matrix, which is worth its extra row-statistics pass only once that matrix is large.", "priority": 43, "when": ["materializedF32Ok", "materializedWideSimdOk", "biasOk", "materializedGemmFusedSoftmaxWorthIt"], "constants": { "qNumHeads": "attrs.num_heads", "kvNumHeads": "attrs.num_heads", "headDim": "headDim", "materializedInnerTile": "tunables.MATERIALIZED_INNER_TILE", "materializedQueryTile": "tunables.MATERIALIZED_QUERY_TILE", "materializedKeyTile": "tunables.MATERIALIZED_KEY_TILE", "materializedWorkgroupDim": "tunables.MATERIALIZED_WORKGROUP_DIM", "applyTileN": "materializedApplyTileN", "hasBias": true, "useSubgroups": "device.features.has(\"subgroups\")", "statSlots": "materializedGemmStatSlots", "statQuerySeq": "dim(shapes.query, 1)", "materializedRowStatsWg": "materializedRowStatsWg" }, "intermediates": [ { "id": "materializedScores", "dtype": "float32", "shape": "[dim(shapes.query, 0) * attrs.num_heads * dim(shapes.query, 1) * dim(shapes.key, 1)]" }, { "id": "materializedRowStats", "dtype": "float32", "shape": "[materializedRowStatsElements]" }, { "id": "materializedScorePartials", "dtype": "float32", "shape": "[materializedGemmScorePartialElements]" } ], "passes": [ { "id": "scores", "name": "MultiHeadAttention.MaterializedScoresBiasFused", "source": { "shader": "attn-materialized-score-f32.wgsl.jinja", "inputs": { "layout": "\"bsh\"", "emitRowStats": true } }, "bindings": "materializedScoreStatsBias", "dispatch": { "x": "ceilDiv(dim(shapes.key, 1), tunables.MATERIALIZED_KEY_TILE)", "y": "ceilDiv(dim(shapes.query, 1), tunables.MATERIALIZED_QUERY_TILE)", "z": "dim(shapes.query, 0) * attrs.num_heads" } }, { "id": "rowstats", "name": "MultiHeadAttention.MaterializedGemmRowStatsCombineBias", "source": { "shader": "attn-materialized-rowstats-combine-f32.wgsl.jinja", "inputs": { "maxOnly": true } }, "bindings": "materializedRowStatsCombine", "dispatch": { "gridStride": "dim(shapes.query, 0) * attrs.num_heads * dim(shapes.query, 1)", "workgroupSize": "materializedRowStatsWg" } }, { "id": "apply", "name": "MultiHeadAttention.MaterializedApplyBiasFused", "source": { "shader": "attn-materialized-apply-f32.wgsl.jinja", "inputs": { "layout": "\"bsh\"", "fusedSoftmax": true } }, "bindings": "materializedApplyFusedBias", "dispatch": { "x": "ceilDiv(headDim, constants.applyTileN)", "y": "ceilDiv(dim(shapes.query, 1), tunables.MATERIALIZED_QUERY_TILE)", "z": "dim(shapes.query, 0) * attrs.num_heads" } } ] }, { "id": "qkv_no_bias_flash_q32_broadcast", "priority": 30, "requires": { "features": ["subgroups", "shader-f16"] }, "when": ["tensorDtypes.query == \"float16\"", "not present.biasT", "flashShapeOk", "(dim(shapes.query, 2) / attrs.num_heads) % 4 == 0", "(dim(shapes.query, 2) / attrs.num_heads) >= 64", "(dim(shapes.query, 2) / attrs.num_heads) <= 256", "dim(shapes.query, 1) >= 31", "subgroupsWave32"], "constants": { "usesF16": true, "scalar": "\"f16\"", "inputVec4": "\"vec4\"", "inputElement": "\"vec4\"", "outputElement": "\"vec4\"", "qNumHeads": "attrs.num_heads", "kvNumHeads": "attrs.num_heads", "headDim": "dim(shapes.query, 2) / attrs.num_heads", "headDimV4": "(dim(shapes.query, 2) / attrs.num_heads) / 4", "kStep": 64, "qkGroups": 16 }, "passes": [ { "id": "main", "name": "MultiHeadAttention.FlashQ32Broadcast", "source": { "shader": "attn-flash-q32-broadcast.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, "bindings": "noBias", "dispatch": { "x": "ceilDiv(dim(shapes.query, 1), 32)", "y": "attrs.num_heads", "z": "dim(shapes.query, 0)" } } ] }, { "id": "qkv_no_bias_flash_q32_shared", "priority": 29, "requires": { "features": ["shader-f16"] }, "when": ["tensorDtypes.query == \"float16\"", "not present.biasT", "flashShapeOk", "(dim(shapes.query, 2) / attrs.num_heads) % 4 == 0", "(dim(shapes.query, 2) / attrs.num_heads) >= 64", "(dim(shapes.query, 2) / attrs.num_heads) <= 256", "dim(shapes.query, 1) >= 31", "ceilDiv(dim(shapes.query, 1), 64) <= device.limits.maxComputeWorkgroupsPerDimension", "((dim(shapes.query, 2) / attrs.num_heads) / 4) * 32 * 16 <= device.limits.maxComputeWorkgroupStorageSize"], "constants": { "usesF16": true, "scalar": "\"f16\"", "inputVec4": "\"vec4\"", "inputElement": "\"vec4\"", "outputElement": "\"vec4\"", "qNumHeads": "attrs.num_heads", "kvNumHeads": "attrs.num_heads", "headDim": "dim(shapes.query, 2) / attrs.num_heads", "headDimV4": "(dim(shapes.query, 2) / attrs.num_heads) / 4", "kStep": 32, "qkGroups": 8, "qStep": 64 }, "passes": [ { "id": "main", "name": "MultiHeadAttention.FlashQ32Shared", "source": { "shader": "attn-flash-q32-broadcast.wgsl.jinja", "inputs": { "layout": "\"bsh\"", "useSubgroups": "false" } }, "bindings": "noBias", "dispatch": { "x": "ceilDiv(dim(shapes.query, 1), 64)", "y": "attrs.num_heads", "z": "dim(shapes.query, 0)" } } ] }, { "id": "qkv_no_bias_attn_bias", "description": "Portable scalar online-attention fallback for legal shapes not admitted by tiled or flash routes. It avoids score scratch while supporting bias and additive-mask forms.", "priority": 0, "when": "not present.biasT and fallbackMaskShapeOk", "constants": { "headsFromParams": false, "hasBias": false, "hasCausal": true, "hasKeyLimit": false, "hasMask": true, "scaleFallbackRsqrt": true, "usesF16": "tensorDtypes.query == \"float16\"", "inputScalar": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "outputScalar": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "inputElement": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "outputElement": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "maskElement": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "qNumHeads": "attrs.num_heads", "kvNumHeads": "attrs.num_heads", "qHidden": "dim(shapes.query, 2)", "kvHidden": "dim(shapes.query, 2)", "headDim": "dim(shapes.query, 2) / attrs.num_heads", "workgroupSize": "portableWorkgroupSize" }, "passes": [ { "id": "main", "name": "MultiHeadAttention", "source": { "shader": "attn-online-scalar.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, "bindings": "noBiasMask", "dispatch": { "x": "min(dim(shapes.query, 1), device.limits.maxComputeWorkgroupsPerDimension)", "y": "attrs.num_heads", "z": "dim(shapes.query, 0)" } } ] }, { "id": "qkv_bias_attn_bias", "description": "Portable scalar online-attention fallback for legal shapes not admitted by tiled or flash routes. It avoids score scratch while supporting bias and additive-mask forms.", "priority": 0, "when": "biasOk and fallbackMaskShapeOk", "constants": { "headsFromParams": false, "hasBias": true, "hasCausal": true, "hasKeyLimit": false, "hasMask": true, "scaleFallbackRsqrt": true, "usesF16": "tensorDtypes.query == \"float16\"", "inputScalar": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "outputScalar": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "inputElement": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "outputElement": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "maskElement": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "qNumHeads": "attrs.num_heads", "kvNumHeads": "attrs.num_heads", "qHidden": "dim(shapes.query, 2)", "kvHidden": "dim(shapes.query, 2)", "headDim": "dim(shapes.query, 2) / attrs.num_heads", "workgroupSize": "portableWorkgroupSize" }, "passes": [ { "id": "main", "name": "MultiHeadAttention", "source": { "shader": "attn-online-scalar.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, "bindings": "biasMask", "dispatch": { "x": "min(dim(shapes.query, 1), device.limits.maxComputeWorkgroupsPerDimension)", "y": "attrs.num_heads", "z": "dim(shapes.query, 0)" } } ] }, { "id": "qkv_no_bias", "description": "Portable scalar online-attention fallback for legal shapes not admitted by tiled or flash routes. It avoids score scratch while supporting bias and additive-mask forms.", "priority": 0, "when": "not present.biasT and fallbackShapeOk and not flashShapeOk", "constants": { "headsFromParams": false, "hasBias": false, "hasCausal": true, "hasKeyLimit": false, "scaleFallbackRsqrt": true, "usesF16": "tensorDtypes.query == \"float16\"", "inputScalar": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "outputScalar": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "inputElement": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "outputElement": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "qNumHeads": "attrs.num_heads", "kvNumHeads": "attrs.num_heads", "qHidden": "dim(shapes.query, 2)", "kvHidden": "dim(shapes.query, 2)", "headDim": "dim(shapes.query, 2) / attrs.num_heads", "workgroupSize": "portableWorkgroupSize" }, "passes": [ { "id": "main", "name": "MultiHeadAttention", "source": { "shader": "attn-online-scalar.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, "bindings": "noBias", "dispatch": { "x": "min(dim(shapes.query, 1), device.limits.maxComputeWorkgroupsPerDimension)", "y": "attrs.num_heads", "z": "dim(shapes.query, 0)" } } ] }, { "id": "qkv_bias", "description": "Portable scalar online-attention fallback for legal shapes not admitted by tiled or flash routes. It avoids score scratch while supporting bias and additive-mask forms.", "priority": 0, "when": "biasOk and fallbackShapeOk and not flashShapeOk", "constants": { "headsFromParams": false, "hasBias": true, "hasCausal": true, "hasKeyLimit": false, "scaleFallbackRsqrt": true, "usesF16": "tensorDtypes.query == \"float16\"", "inputScalar": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "outputScalar": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "inputElement": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "outputElement": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "qNumHeads": "attrs.num_heads", "kvNumHeads": "attrs.num_heads", "qHidden": "dim(shapes.query, 2)", "kvHidden": "dim(shapes.query, 2)", "headDim": "dim(shapes.query, 2) / attrs.num_heads", "workgroupSize": "portableWorkgroupSize" }, "passes": [ { "id": "main", "name": "MultiHeadAttention", "source": { "shader": "attn-online-scalar.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, "bindings": "bias", "dispatch": { "x": "min(dim(shapes.query, 1), device.limits.maxComputeWorkgroupsPerDimension)", "y": "attrs.num_heads", "z": "dim(shapes.query, 0)" } } ] }, { "id": "qkv_no_bias_flash", "description": "Computes one-pass online-softmax attention without score scratch. Family suffixes select subgroup, portable, and additive-attention-bias forms without changing the algorithm.", "priority": 20, "when": ["device.features.has(\"subgroups\")", "not present.biasT", "flashShapeOk"], "constants": { "headsFromParams": false, "hasBias": false, "hasCausal": true, "combineSubgroups": true, "hasMask": false, "maskIsBool": false, "maskElement": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "usesF16": "tensorDtypes.query == \"float16\"", "scalar": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "inputVec4": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "inputElement": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "outputElement": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "qNumHeads": "attrs.num_heads", "kvNumHeads": "attrs.num_heads", "qHiddenV4": "dim(shapes.query, 2) / 4", "kvHiddenV4": "dim(shapes.query, 2) / 4", "headDim": "dim(shapes.query, 2) / attrs.num_heads", "headDimV4": "(dim(shapes.query, 2) / attrs.num_heads) / 4" }, "passes": [ { "id": "main", "name": "MultiHeadAttention.Flash", "source": { "shader": "attn-flash-online.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, "bindings": "noBias", "dispatch": { "x": "dim(shapes.query, 1)", "y": "attrs.num_heads", "z": "dim(shapes.query, 0)" } } ] }, { "id": "qkv_bias_flash", "description": "Computes one-pass online-softmax attention with fused projection bias and no score scratch. Family suffixes select subgroup, portable, and additive-attention-bias forms without changing the algorithm.", "priority": 20, "when": ["device.features.has(\"subgroups\")", "biasOk", "flashShapeOk"], "constants": { "headsFromParams": false, "hasBias": true, "hasCausal": true, "combineSubgroups": true, "hasMask": false, "maskIsBool": false, "maskElement": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "usesF16": "tensorDtypes.query == \"float16\"", "scalar": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "inputVec4": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "inputElement": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "outputElement": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "inputScalar": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "qNumHeads": "attrs.num_heads", "kvNumHeads": "attrs.num_heads", "qHidden": "dim(shapes.query, 2)", "qHiddenV4": "dim(shapes.query, 2) / 4", "kvHiddenV4": "dim(shapes.query, 2) / 4", "headDim": "dim(shapes.query, 2) / attrs.num_heads", "headDimV4": "(dim(shapes.query, 2) / attrs.num_heads) / 4" }, "passes": [ { "id": "main", "name": "MultiHeadAttention.Flash", "source": { "shader": "attn-flash-online.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, "bindings": "bias", "dispatch": { "x": "dim(shapes.query, 1)", "y": "attrs.num_heads", "z": "dim(shapes.query, 0)" } } ] }, { "id": "qkv_no_bias_flash_nosg", "description": "Computes one-pass online-softmax attention without score scratch. Family suffixes select subgroup, portable, and additive-attention-bias forms without changing the algorithm.", "priority": 18, "when": ["true", "not present.biasT", "flashShapeOk"], "constants": { "headsFromParams": false, "hasBias": false, "hasCausal": true, "combineSubgroups": false, "hasMask": false, "maskIsBool": false, "maskElement": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "usesF16": "tensorDtypes.query == \"float16\"", "scalar": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "inputVec4": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "inputElement": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "outputElement": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "qNumHeads": "attrs.num_heads", "kvNumHeads": "attrs.num_heads", "qHiddenV4": "dim(shapes.query, 2) / 4", "kvHiddenV4": "dim(shapes.query, 2) / 4", "headDim": "dim(shapes.query, 2) / attrs.num_heads", "headDimV4": "(dim(shapes.query, 2) / attrs.num_heads) / 4" }, "passes": [ { "id": "main", "name": "MultiHeadAttention.NoBiasOnlineFlashNoSg", "source": { "shader": "attn-flash-online.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, "bindings": "noBias", "dispatch": { "x": "dim(shapes.query, 1)", "y": "attrs.num_heads", "z": "dim(shapes.query, 0)" } } ] }, { "id": "qkv_bias_flash_nosg", "description": "Computes one-pass online-softmax attention with fused projection bias and no score scratch. Family suffixes select subgroup, portable, and additive-attention-bias forms without changing the algorithm.", "priority": 17, "when": ["true", "biasOk", "flashShapeOk"], "constants": { "headsFromParams": false, "hasBias": true, "hasCausal": true, "combineSubgroups": false, "hasMask": false, "maskIsBool": false, "maskElement": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "usesF16": "tensorDtypes.query == \"float16\"", "scalar": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "inputVec4": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "inputElement": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "outputElement": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "inputScalar": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "qNumHeads": "attrs.num_heads", "kvNumHeads": "attrs.num_heads", "qHidden": "dim(shapes.query, 2)", "qHiddenV4": "dim(shapes.query, 2) / 4", "kvHiddenV4": "dim(shapes.query, 2) / 4", "headDim": "dim(shapes.query, 2) / attrs.num_heads", "headDimV4": "(dim(shapes.query, 2) / attrs.num_heads) / 4" }, "passes": [ { "id": "main", "name": "MultiHeadAttention.BiasOnlineFlashNoSg", "source": { "shader": "attn-flash-online.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, "bindings": "bias", "dispatch": { "x": "dim(shapes.query, 1)", "y": "attrs.num_heads", "z": "dim(shapes.query, 0)" } } ] }, { "id": "qkv_no_bias_flash_attn_bias", "description": "Computes one-pass online-softmax attention without score scratch. Family suffixes select subgroup, portable, and additive-attention-bias forms without changing the algorithm.", "priority": 20, "when": ["device.features.has(\"subgroups\")", "not present.biasT", "flashMaskShapeOk"], "constants": { "headsFromParams": false, "hasBias": false, "hasCausal": true, "combineSubgroups": true, "hasMask": true, "maskIsBool": false, "maskElement": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "usesF16": "tensorDtypes.query == \"float16\"", "scalar": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "inputVec4": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "inputElement": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "outputElement": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "qNumHeads": "attrs.num_heads", "kvNumHeads": "attrs.num_heads", "qHiddenV4": "dim(shapes.query, 2) / 4", "kvHiddenV4": "dim(shapes.query, 2) / 4", "headDim": "dim(shapes.query, 2) / attrs.num_heads", "headDimV4": "(dim(shapes.query, 2) / attrs.num_heads) / 4" }, "passes": [ { "id": "main", "name": "MultiHeadAttention.Flash", "source": { "shader": "attn-flash-online.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, "bindings": "noBiasMask", "dispatch": { "x": "dim(shapes.query, 1)", "y": "attrs.num_heads", "z": "dim(shapes.query, 0)" } } ] }, { "id": "qkv_bias_flash_attn_bias", "description": "Computes one-pass online-softmax attention with fused projection bias and no score scratch. Family suffixes select subgroup, portable, and additive-attention-bias forms without changing the algorithm.", "priority": 20, "when": ["device.features.has(\"subgroups\")", "biasOk", "flashMaskShapeOk"], "constants": { "headsFromParams": false, "hasBias": true, "hasCausal": true, "combineSubgroups": true, "hasMask": true, "maskIsBool": false, "maskElement": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "usesF16": "tensorDtypes.query == \"float16\"", "scalar": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "inputVec4": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "inputElement": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "outputElement": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "inputScalar": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "qNumHeads": "attrs.num_heads", "kvNumHeads": "attrs.num_heads", "qHidden": "dim(shapes.query, 2)", "qHiddenV4": "dim(shapes.query, 2) / 4", "kvHiddenV4": "dim(shapes.query, 2) / 4", "headDim": "dim(shapes.query, 2) / attrs.num_heads", "headDimV4": "(dim(shapes.query, 2) / attrs.num_heads) / 4" }, "passes": [ { "id": "main", "name": "MultiHeadAttention.Flash", "source": { "shader": "attn-flash-online.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, "bindings": "biasMask", "dispatch": { "x": "dim(shapes.query, 1)", "y": "attrs.num_heads", "z": "dim(shapes.query, 0)" } } ] }, { "id": "qkv_no_bias_flash_attn_bias_nosg", "description": "Computes one-pass online-softmax attention without score scratch. Family suffixes select subgroup, portable, and additive-attention-bias forms without changing the algorithm.", "priority": 18, "when": ["true", "not present.biasT", "flashMaskShapeOk"], "constants": { "headsFromParams": false, "hasBias": false, "hasCausal": true, "combineSubgroups": false, "hasMask": true, "maskIsBool": false, "maskElement": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "usesF16": "tensorDtypes.query == \"float16\"", "scalar": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "inputVec4": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "inputElement": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "outputElement": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "qNumHeads": "attrs.num_heads", "kvNumHeads": "attrs.num_heads", "qHiddenV4": "dim(shapes.query, 2) / 4", "kvHiddenV4": "dim(shapes.query, 2) / 4", "headDim": "dim(shapes.query, 2) / attrs.num_heads", "headDimV4": "(dim(shapes.query, 2) / attrs.num_heads) / 4" }, "passes": [ { "id": "main", "name": "MultiHeadAttention.NoBiasOnlineFlashNoSg", "source": { "shader": "attn-flash-online.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, "bindings": "noBiasMask", "dispatch": { "x": "dim(shapes.query, 1)", "y": "attrs.num_heads", "z": "dim(shapes.query, 0)" } } ] }, { "id": "qkv_bias_flash_attn_bias_nosg", "description": "Computes one-pass online-softmax attention with fused projection bias and no score scratch. Family suffixes select subgroup, portable, and additive-attention-bias forms without changing the algorithm.", "priority": 17, "when": ["true", "biasOk", "flashMaskShapeOk"], "constants": { "headsFromParams": false, "hasBias": true, "hasCausal": true, "combineSubgroups": false, "hasMask": true, "maskIsBool": false, "maskElement": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "usesF16": "tensorDtypes.query == \"float16\"", "scalar": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "inputVec4": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "inputElement": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "outputElement": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", "inputScalar": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", "qNumHeads": "attrs.num_heads", "kvNumHeads": "attrs.num_heads", "qHidden": "dim(shapes.query, 2)", "qHiddenV4": "dim(shapes.query, 2) / 4", "kvHiddenV4": "dim(shapes.query, 2) / 4", "headDim": "dim(shapes.query, 2) / attrs.num_heads", "headDimV4": "(dim(shapes.query, 2) / attrs.num_heads) / 4" }, "passes": [ { "id": "main", "name": "MultiHeadAttention.BiasOnlineFlashNoSg", "source": { "shader": "attn-flash-online.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, "bindings": "biasMask", "dispatch": { "x": "dim(shapes.query, 1)", "y": "attrs.num_heads", "z": "dim(shapes.query, 0)" } } ] }, { "id": "qkv_no_bias_materialized_sgmat_f32", "description": "Materializes FP32 scores so subgroup-matrix operations can compute the score and value products, normalizing them in a separate softmax pass. Selected when the score matrix is small enough to stay cache-resident, where that pass costs less than the staging the fused route forces on apply.", "priority": 52, "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] }, "when": ["materializedSgmatOk", "not present.biasT", "not materializedFusedSoftmaxWorthIt"], "constants": { "qNumHeads": "attrs.num_heads", "headDim": "headDim", "qHidden": "dim(shapes.query, 2)", "materializedSgmatQueryTile": "materializedSgmatQueryTile", "materializedSgmatKeyTile": "materializedSgmatKeyTile", "materializedSgmatInnerTile": "materializedSgmatInnerTile", "materializedSgmatDirectScoreStore": "materializedSgmatDirectScoreStore", "materializedSgmatDirectApplyStore": "materializedSgmatDirectApplyStore", "materializedSgmatRuntimeDirectStore": "materializedSgmatRuntimeDirectStore", "materializedSoftmaxWg": "materializedCachedSoftmaxWg if materializedCachedSoftmaxOk else tunables.MATERIALIZED_SOFTMAX_WORKGROUP_SIZE", "materializedSoftmaxCols": "dim(shapes.key, 1)", "materializedSoftmaxCols4": "dim(shapes.key, 1) / 4", "softmaxElementType": "\"vec4\" if materializedCachedSoftmaxOk else \"f32\"", "hasBias": false, "useSubgroups": true }, "intermediates": [ { "id": "materializedScores", "dtype": "float32", "shape": "[dim(shapes.query, 0) * attrs.num_heads * dim(shapes.query, 1) * dim(shapes.key, 1)]" } ], "passes": [ { "id": "scores", "name": "MultiHeadAttention.MaterializedScoresSgmat", "source": { "shader": "attn-materialized-sgmat-f32.wgsl.jinja", "inputs": { "phase": "\"score\"" } }, "bindings": "materializedScore", "dispatch": { "x": "ceilDiv(dim(shapes.key, 1), materializedSgmatKeyTile)", "y": "ceilDiv(dim(shapes.query, 1), materializedSgmatQueryTile)", "z": "dim(shapes.query, 0) * attrs.num_heads" } }, { "id": "softmax", "name": "MultiHeadAttention.MaterializedSoftmaxSgmat", "source": { "shader": "attn-materialized-softmax-f32.wgsl.jinja", "inputs": { "cacheVec4": "materializedCachedSoftmaxOk" } }, "bindings": "materializedSoftmax", "dispatch": { "workgroups": "dim(shapes.query, 0) * attrs.num_heads * dim(shapes.query, 1)" } }, { "id": "apply", "name": "MultiHeadAttention.MaterializedApplySgmat", "source": { "shader": "attn-materialized-sgmat-f32.wgsl.jinja", "inputs": { "phase": "\"apply\"" } }, "bindings": "materializedApply", "dispatch": { "x": "ceilDiv(headDim, materializedSgmatKeyTile)", "y": "ceilDiv(dim(shapes.query, 1), materializedSgmatQueryTile)", "z": "dim(shapes.query, 0) * attrs.num_heads" } } ] }, { "id": "qkv_bias_materialized_sgmat_f32", "description": "Materializes FP32 scores so subgroup-matrix operations can compute the score and value products, normalizing them in a separate softmax pass. Selected when the score matrix is small enough to stay cache-resident, where that pass costs less than the staging the fused route forces on apply.", "priority": 52, "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] }, "when": ["materializedSgmatOk", "biasOk", "not materializedFusedSoftmaxWorthIt"], "constants": { "qNumHeads": "attrs.num_heads", "headDim": "headDim", "qHidden": "dim(shapes.query, 2)", "materializedSgmatQueryTile": "materializedSgmatQueryTile", "materializedSgmatKeyTile": "materializedSgmatKeyTile", "materializedSgmatInnerTile": "materializedSgmatInnerTile", "materializedSgmatDirectScoreStore": "materializedSgmatDirectScoreStore", "materializedSgmatDirectApplyStore": "materializedSgmatDirectApplyStore", "materializedSgmatRuntimeDirectStore": "materializedSgmatRuntimeDirectStore", "materializedSoftmaxWg": "materializedCachedSoftmaxWg if materializedCachedSoftmaxOk else tunables.MATERIALIZED_SOFTMAX_WORKGROUP_SIZE", "materializedSoftmaxCols": "dim(shapes.key, 1)", "materializedSoftmaxCols4": "dim(shapes.key, 1) / 4", "softmaxElementType": "\"vec4\" if materializedCachedSoftmaxOk else \"f32\"", "hasBias": true, "useSubgroups": true }, "intermediates": [ { "id": "materializedScores", "dtype": "float32", "shape": "[dim(shapes.query, 0) * attrs.num_heads * dim(shapes.query, 1) * dim(shapes.key, 1)]" } ], "passes": [ { "id": "scores", "name": "MultiHeadAttention.MaterializedScoresSgmatBias", "source": { "shader": "attn-materialized-sgmat-f32.wgsl.jinja", "inputs": { "phase": "\"score\"" } }, "bindings": "materializedScoreBias", "dispatch": { "x": "ceilDiv(dim(shapes.key, 1), materializedSgmatKeyTile)", "y": "ceilDiv(dim(shapes.query, 1), materializedSgmatQueryTile)", "z": "dim(shapes.query, 0) * attrs.num_heads" } }, { "id": "softmax", "name": "MultiHeadAttention.MaterializedSoftmaxSgmatBias", "source": { "shader": "attn-materialized-softmax-f32.wgsl.jinja", "inputs": { "cacheVec4": "materializedCachedSoftmaxOk" } }, "bindings": "materializedSoftmax", "dispatch": { "workgroups": "dim(shapes.query, 0) * attrs.num_heads * dim(shapes.query, 1)" } }, { "id": "apply", "name": "MultiHeadAttention.MaterializedApplySgmatBias", "source": { "shader": "attn-materialized-sgmat-f32.wgsl.jinja", "inputs": { "phase": "\"apply\"" } }, "bindings": "materializedApplyBias", "dispatch": { "x": "ceilDiv(headDim, materializedSgmatKeyTile)", "y": "ceilDiv(dim(shapes.query, 1), materializedSgmatQueryTile)", "z": "dim(shapes.query, 0) * attrs.num_heads" } } ] } ] }