diff --git "a/build/webgpu/manifest.json" "b/build/webgpu/manifest.json" new file mode 100644--- /dev/null +++ "b/build/webgpu/manifest.json" @@ -0,0 +1,5197 @@ +{ + "domain": "com.microsoft", + "name": "GroupQueryAttention", + "sinceVersion": 1, + "description": "Grouped-query attention for explicit BSH Q/K/V and BNSH caches. Direct Q/K/V supports bidirectional attention or causal local windows and may store its generated float cache independently as float16 or float32; existing unquantized cache inputs match the Q/K/V dtype. Causal cache paths support rotary embeddings, sliding windows, bias, head sinks, softcap, smooth softmax, and paired Q/K RMS normalization. Int8/int4 caches require float32 Q/K/V and output; int4 is prompt-only. Packed QKV, position IDs, interleaved rotary, bfloat16/float8, and diagnostic QK output are not implemented.", + "inputs": [ + { + "role": "query", + "dtype": "T", + "rank": 3, + "description": "Query tensor of shape `(batch_size, sequence_length, num_heads * head_size)`." + }, + { + "role": "key", + "dtype": "T", + "rank": 3, + "description": "Key tensor of shape `(batch_size, kv_sequence_length, kv_num_heads * head_size)`." + }, + { + "role": "value", + "dtype": "T", + "rank": 3, + "description": "Value tensor of shape `(batch_size, kv_sequence_length, kv_num_heads * head_size)`." + }, + { + "role": "past_key", + "dtype": "T_CACHE", + "rank": 4, + "optional": true, + "description": "Optional cached key state in BNSH format. Its sequence axis is `max_sequence_length` when the past and present buffers are shared, otherwise `past_sequence_length`; int4 stores each signed value as a +8-biased nibble, with the even head coordinate low, packing two values per logical uint8 element and widening each byte to one u32 WebGPU buffer word." + }, + { + "role": "past_value", + "dtype": "T_CACHE", + "rank": 4, + "optional": true, + "description": "Optional cached value state in BNSH format with the same length and packing semantics as `past_key`." + }, + { + "role": "seqlens_k", + "dtype": "M", + "rank": 1, + "description": "Int32 tensor of shape `(batch_size)` containing each sample's total sequence length minus one." + }, + { + "role": "total_sequence_length", + "dtype": "M", + "rank": 1, + "description": "Length-one int32 tensor containing the maximum total sequence length (past plus new) in the batch." + }, + { + "role": "cos_cache", + "dtype": "T", + "rank": 2, + "optional": true, + "description": "Optional cosine cache for rotary embeddings with shape `(max_sequence_length, head_size / 2)`." + }, + { + "role": "sin_cache", + "dtype": "T", + "rank": 2, + "optional": true, + "description": "Optional sine cache for rotary embeddings with shape `(max_sequence_length, head_size / 2)`." + }, + { + "role": "attention_bias", + "dtype": "T", + "rank": 4, + "optional": true, + "description": "Optional additive term for QK scores with shape `(batch_size or 1, num_heads or 1, sequence_length, total_sequence_length)`; the first two dimensions broadcast." + }, + { + "role": "head_sink", + "dtype": "T", + "rank": 1, + "optional": true, + "description": "Optional per-head smooth factor of shape `(num_heads)` added to the softmax denominator." + }, + { + "role": "k_scale", + "dtype": "T_KV_SCALE", + "rank": 1, + "optional": true, + "description": "Optional float32 key-cache scale: one value for `PER_TENSOR`, or `kv_num_heads * head_size` values for `PER_CHANNEL`." + }, + { + "role": "v_scale", + "dtype": "T_KV_SCALE", + "rank": 1, + "optional": true, + "description": "Optional float32 value-cache scale with the same shape convention as `k_scale`." + }, + { + "role": "q_norm_weight", + "dtype": "T", + "rank": 1, + "optional": true, + "description": "Optional per-head RMS-normalization weight of shape `(head_size)` applied to queries before rotary embedding. It must be provided together with `k_norm_weight`." + }, + { + "role": "k_norm_weight", + "dtype": "T", + "rank": 1, + "optional": true, + "description": "Optional per-head RMS-normalization weight of shape `(head_size)` applied to keys before rotary embedding. It must be provided together with `q_norm_weight`." + } + ], + "outputs": [ + { + "role": "output", + "dtype": "T", + "rank": 3, + "shape": "shapes.query", + "description": "Attention output of shape `(batch_size, sequence_length, hidden_size)`." + }, + { + "role": "present_key", + "dtype": "T_CACHE", + "rank": 4, + "description": "Updated key cache in BNSH format. Its sequence axis is `max_sequence_length` for a shared buffer, otherwise `past_sequence_length + kv_sequence_length`; int4 stores each signed value as a +8-biased nibble, with the even head coordinate low, packing two values per logical uint8 element and widening each byte to one u32 WebGPU buffer word." + }, + { + "role": "present_value", + "dtype": "T_CACHE", + "rank": 4, + "description": "Updated value cache in BNSH format with the same length and packing semantics as `present_key`." + } + ], + "attributes": { + "causal": 1, + "do_rotary": 0, + "k_quant_type": "NONE", + "local_window_size": -1, + "qk_norm_epsilon": 0.000001, + "sliding_window_cache": 0, + "smooth_softmax": -1, + "softcap": 0, + "v_quant_type": "NONE" + }, + "attributeDescriptions": { + "causal": "Whether to apply a causal mask. Set to 0 for bidirectional attention; `local_window_size` must then be -1.", + "do_rotary": "Set to 1 to apply rotary position embeddings. The default 0 disables them.", + "k_quant_type": "Key-cache quantization mode: `NONE`, `PER_TENSOR`, or `PER_CHANNEL`.", + "kv_cache_bit_width": "Quantized cache bit width, either 8 or 4. Four-bit values are packed two per uint8 element.", + "kv_num_heads": "Number of key/value attention heads.", + "local_window_size": "Left window size for causal local attention. The default -1 disables local attention, and the value must be -1 when `causal` is 0.", + "num_heads": "Number of query attention heads.", + "qk_norm_epsilon": "Epsilon for the per-head Q/K RMS normalization applied when both normalization weights are provided.", + "scale": "Optional QK score scale; zero or omission selects `1 / sqrt(head_size)`.", + "sliding_window_cache": "Set to 1 when past/present caches are fixed-size window buffers that evict old tokens from the front. Requires `local_window_size > 0` and enough cache capacity.", + "smooth_softmax": "Set to 1 to enable the smooth-softmax denominator term.", + "softcap": "Positive softcap applied to attention scores. The default 0 disables soft-capping.", + "v_quant_type": "Value-cache quantization mode: `NONE`, `PER_TENSOR`, or `PER_CHANNEL`." + }, + "attributeConstraints": { + "causal": { "values": [0, 1] }, + "do_rotary": { "values": [0, 1] }, + "k_quant_type": { "values": ["NONE", "PER_TENSOR", "PER_CHANNEL"] }, + "kv_cache_bit_width": { "values": [4, 8] }, + "kv_num_heads": { "required": true }, + "num_heads": { "required": true }, + "sliding_window_cache": { "values": [0, 1] }, + "v_quant_type": { "values": ["NONE", "PER_TENSOR", "PER_CHANNEL"] } + }, + "typeConstraints": { + "T": ["float32", "float16"], + "T_CACHE": ["float32", "float16", "uint8", "int8"], + "T_KV_SCALE": ["float32"], + "M": ["int32"] + }, + "args": { + "queryT": { "kind": "tensor", "semantic": "query", "role": "input" }, + "keyT": { "kind": "tensor", "semantic": "key", "role": "input" }, + "valueT": { "kind": "tensor", "semantic": "value", "role": "input" }, + "pastKeyT": { "kind": "tensor", "semantic": "past_key", "role": "input", "required": false }, + "pastValueT": { "kind": "tensor", "semantic": "past_value", "role": "input", "required": false }, + "seqlensKT": { "kind": "tensor", "semantic": "seqlens_k", "role": "input" }, + "totalSequenceLengthT": { "kind": "tensor", "semantic": "total_sequence_length", "role": "input" }, + "cosCacheT": { "kind": "tensor", "semantic": "cos_cache", "role": "input", "required": false }, + "sinCacheT": { "kind": "tensor", "semantic": "sin_cache", "role": "input", "required": false }, + "attentionBiasT": { "kind": "tensor", "semantic": "attention_bias", "role": "input", "required": false }, + "headSinkT": { "kind": "tensor", "semantic": "head_sink", "role": "input", "required": false }, + "kScaleT": { "kind": "tensor", "semantic": "k_scale", "role": "input", "required": false }, + "vScaleT": { "kind": "tensor", "semantic": "v_scale", "role": "input", "required": false }, + "qNormWeightT": { "kind": "tensor", "semantic": "q_norm_weight", "role": "input", "required": false }, + "kNormWeightT": { "kind": "tensor", "semantic": "k_norm_weight", "role": "input", "required": false }, + "outputT": { "kind": "tensor", "semantic": "output", "role": "output" }, + "presentKeyT": { "kind": "tensor", "semantic": "present_key", "role": "output" }, + "presentValueT": { "kind": "tensor", "semantic": "present_value", "role": "output" } + }, + "tunables": { + "WORKGROUP_SIZE": 256, + "COPY_WORKGROUP_SIZE": 64, + "SCALAR_WORKGROUP_SIZE": 64, + "COOPERATIVE_QUERY_THRESHOLD": 6144, + "MAX_SPLITS": 16, + "QKV_SPLIT_TILE_K": 128, + "CACHED_SPLIT_TILE_K": 256, + "FLASH_MIN_HEAD_DIM": 32, + "FLASH_MAX_HEAD_DIM": 256, + "PREFILL_QUERY_TILE": 32, + "PREFILL_LANES_PER_QUERY": 4, + "PREFILL_QUERY_TILE_F16": 16, + "PREFILL_LANES_PER_QUERY_F16": 8, + "QKV_PREFILL_MIN_QUERY_TOKENS": 31, + "QKV_DECODE_MIN_KV_TOKENS": 512, + "CACHED_FLASH_MIN_HEAD_DIM": 64, + "CACHED_DECODE_MIN_KV_TOKENS": 1024, + "CLUSTER_TILE_K_F32": 8, + "CLUSTER_TILE_K_F16": 8, + "NO_SG_TILE_K_MAX": 16, + "COOPERATIVE_WORKGROUP_SIZE": 32, + "CLUSTER_MAX_SLICE": 8 + }, + "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", + "windowCacheRequested": "attrs.sliding_window_cache == 1", + "headDim": "dim(shapes.query, 2) / attrs.num_heads if (ranks.query == 3 and attrs.num_heads > 0) else 0", + "copyWorkgroupSize": "min(tunables.COPY_WORKGROUP_SIZE, min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX))", + "scalarWorkgroupSize": "min(tunables.SCALAR_WORKGROUP_SIZE, min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX))", + "copyWorkgroupOk": "copyWorkgroupSize >= 1", + "scalarWorkgroupOk": "scalarWorkgroupSize >= 1", + "metadataInputsOk": "present.seqlensKT and present.totalSequenceLengthT", + "unquantizedCacheAttrs": "attrs.k_quant_type == \"NONE\" and attrs.v_quant_type == \"NONE\" and not attrs.kv_cache_bit_width and not present.kScaleT and not present.vScaleT", + "qkNormInputsOk": "present.qNormWeightT == present.kNormWeightT", + "rotaryRequested": "attrs.do_rotary == 1", + "noAuxAttentionInputs": "not present.cosCacheT and not rotaryRequested and not present.qNormWeightT and not present.attentionBiasT and not present.headSinkT", + "standardSoftmax": "attrs.softcap == 0 and attrs.smooth_softmax != 1", + "plainAttentionOptions": "noAuxAttentionInputs and standardSoftmax", + "qkvMaskModeOk": "(attrs.causal == 0 and attrs.local_window_size == -1) or (attrs.causal == 1 and attrs.local_window_size > 0)", + "qkvContractOk": "plainAttentionOptions and not windowCacheRequested and copyWorkgroupOk and metadataInputsOk and unquantizedCacheAttrs and qkNormInputsOk and qkvMaskModeOk and (not present.pastKeyT) and (not present.pastValueT) and ranks.query == 3 and ranks.key == 3 and ranks.value == 3 and ranks.outputT == 3 and ranks.presentKeyT == 4 and ranks.presentValueT == 4 and (tensorDtypes.query == \"float32\" or tensorDtypes.query == \"float16\") and f16Ok(tensorDtypes.query) and tensorDtypes.key == tensorDtypes.query and tensorDtypes.value == tensorDtypes.query and tensorDtypes.outputT == tensorDtypes.query and (tensorDtypes.presentKeyT == \"float32\" or tensorDtypes.presentKeyT == \"float16\") and f16Ok(tensorDtypes.presentKeyT) and tensorDtypes.presentValueT == tensorDtypes.presentKeyT and attrs.num_heads > 0 and attrs.kv_num_heads > 0 and attrs.num_heads % attrs.kv_num_heads == 0 and dim(shapes.query, 2) % attrs.num_heads == 0 and dim(shapes.key, 2) == attrs.kv_num_heads * headDim and dim(shapes.value, 2) == dim(shapes.key, 2) and dim(shapes.query, 0) == dim(shapes.key, 0) and dim(shapes.query, 0) == dim(shapes.value, 0) and dim(shapes.key, 1) == dim(shapes.value, 1) and dim(shapes.outputT, 0) == dim(shapes.query, 0) and dim(shapes.outputT, 1) == dim(shapes.query, 1) and dim(shapes.outputT, 2) == dim(shapes.query, 2) and dim(shapes.presentKeyT, 0) == dim(shapes.query, 0) and dim(shapes.presentKeyT, 1) == attrs.kv_num_heads and dim(shapes.presentKeyT, 2) == dim(shapes.key, 1) and dim(shapes.presentKeyT, 3) == headDim and dim(shapes.presentValueT, 0) == dim(shapes.presentKeyT, 0) and dim(shapes.presentValueT, 1) == dim(shapes.presentKeyT, 1) and dim(shapes.presentValueT, 2) == dim(shapes.presentKeyT, 2) and dim(shapes.presentValueT, 3) == dim(shapes.presentKeyT, 3)", + "flashWorkgroupOk": "tunables.WORKGROUP_SIZE <= deviceWorkgroupCap", + "flashShapeOk": "qkvContractOk and flashWorkgroupOk and headDim % 4 == 0 and headDim >= tunables.FLASH_MIN_HEAD_DIM and headDim <= tunables.FLASH_MAX_HEAD_DIM and (dim(shapes.query, 1) * attrs.num_heads >= 248 or (dim(shapes.query, 1) == 1 and dim(shapes.key, 1) >= tunables.QKV_DECODE_MIN_KV_TOKENS) or (dim(shapes.query, 1) > 1 and dim(shapes.key, 1) >= 2048)) and dim(shapes.query, 0) <= device.limits.maxComputeWorkgroupsPerDimension and dim(shapes.query, 1) <= device.limits.maxComputeWorkgroupsPerDimension and attrs.num_heads <= device.limits.maxComputeWorkgroupsPerDimension", + "decodeSplitKShapeOk": "qkvContractOk and flashWorkgroupOk and dim(shapes.query, 1) == 1 and dim(shapes.key, 1) >= tunables.QKV_DECODE_MIN_KV_TOKENS and headDim % 4 == 0 and headDim >= tunables.FLASH_MIN_HEAD_DIM and headDim <= tunables.FLASH_MAX_HEAD_DIM and dim(shapes.query, 0) <= device.limits.maxComputeWorkgroupsPerDimension and attrs.num_heads <= device.limits.maxComputeWorkgroupsPerDimension", + "fallbackShapeOk": "qkvContractOk and headDim <= deviceWorkgroupCap and headDim <= tunables.FLASH_MAX_HEAD_DIM", + "gqaHeadDim": "dim(shapes.queryT, 2) / attrs.num_heads", + "qkvScalar": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", + "qkvInputVec4": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", + "qkvCacheScalar": "\"f16\" if tensorDtypes.presentKeyT == \"float16\" else \"f32\"", + "qkvCacheVec4": "\"vec4\" if tensorDtypes.presentKeyT == \"float16\" else \"vec4\"", + "qkvUsesF16": "tensorDtypes.query == \"float16\" or tensorDtypes.presentKeyT == \"float16\"", + "gqaScalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"", + "gqaInputVec4": "\"vec4\" if tensorDtypes.queryT == \"float16\" else \"vec4\"", + "qkvNumSplits": "min(tunables.MAX_SPLITS, ceilDiv(dim(shapes.key, 1), tunables.QKV_SPLIT_TILE_K))", + "cachedNumSplits": "min(tunables.MAX_SPLITS, ceilDiv(dim(shapes.presentKeyT, 2), tunables.CACHED_SPLIT_TILE_K))", + "cachedBaseOk": "attrs.causal == 1 and metadataInputsOk and qkNormInputsOk", + "sharedKvCacheOk": "cachedBaseOk and unquantizedCacheAttrs and not windowCacheRequested and copyWorkgroupOk and present.pastKeyT and present.pastValueT and dim(shapes.keyT, 1) == 0 and dim(shapes.queryT, 2) % attrs.num_heads == 0 and dim(shapes.pastKeyT, 1) == attrs.kv_num_heads and dim(shapes.pastKeyT, 3) == gqaHeadDim", + "gqaFlashHd": "gqaHeadDim >= tunables.CACHED_FLASH_MIN_HEAD_DIM and gqaHeadDim <= tunables.FLASH_MAX_HEAD_DIM", + "f16OrF32": "(tensorDtypes.queryT == \"float32\" or tensorDtypes.queryT == \"float16\") and f16Ok(tensorDtypes.queryT) and tensorDtypes.keyT == tensorDtypes.queryT and tensorDtypes.valueT == tensorDtypes.queryT and tensorDtypes.outputT == tensorDtypes.queryT and tensorDtypes.presentKeyT == tensorDtypes.queryT and tensorDtypes.presentValueT == tensorDtypes.queryT and ((not present.pastKeyT) or tensorDtypes.pastKeyT == tensorDtypes.queryT) and ((not present.pastValueT) or tensorDtypes.pastValueT == tensorDtypes.queryT)", + "subgroupCluster8": "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", + "qkvPrefillQueryTile": "tunables.PREFILL_QUERY_TILE_F16 if tensorDtypes.query == \"float16\" else tunables.PREFILL_QUERY_TILE", + "qkvPrefillLanesPerQuery": "max(tunables.PREFILL_LANES_PER_QUERY_F16 if tensorDtypes.query == \"float16\" else tunables.PREFILL_LANES_PER_QUERY, ceilDiv(headDim / 4, tunables.CLUSTER_MAX_SLICE))", + "cachedPrefillQueryTile": "tunables.PREFILL_QUERY_TILE_F16 if tensorDtypes.queryT == \"float16\" else tunables.PREFILL_QUERY_TILE", + "cachedPrefillLanesPerQuery": "max(tunables.PREFILL_LANES_PER_QUERY_F16 if tensorDtypes.queryT == \"float16\" else tunables.PREFILL_LANES_PER_QUERY, ceilDiv(gqaHeadDim / 4, tunables.CLUSTER_MAX_SLICE))", + "subgroupQkvPrefillCluster": "not narrowSubgroupRange and device.features.has(\"subgroups\") and has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize >= qkvPrefillLanesPerQuery and device.adapterInfo.subgroupMinSize % qkvPrefillLanesPerQuery == 0 and device.adapterInfo.subgroupMaxSize % qkvPrefillLanesPerQuery == 0", + "subgroupCachedPrefillCluster": "not narrowSubgroupRange and device.features.has(\"subgroups\") and has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize >= cachedPrefillLanesPerQuery and device.adapterInfo.subgroupMinSize % cachedPrefillLanesPerQuery == 0 and device.adapterInfo.subgroupMaxSize % cachedPrefillLanesPerQuery == 0", + "subgroupSplitK": "subgroupCluster8 and device.adapterInfo.subgroupMaxSize <= tunables.WORKGROUP_SIZE and tunables.WORKGROUP_SIZE % device.adapterInfo.subgroupMinSize == 0 and tunables.WORKGROUP_SIZE % device.adapterInfo.subgroupMaxSize == 0", + "gqaQueryCount": "dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)", + "cooperativeWorkgroupOk": "tunables.COOPERATIVE_WORKGROUP_SIZE <= deviceWorkgroupCap", + "gqaCoopWorkgroupBytes": "((tunables.COOPERATIVE_WORKGROUP_SIZE + 1) * gqaHeadDim + tunables.COOPERATIVE_WORKGROUP_SIZE * 2) * 4", + "gqaCoop": "cooperativeWorkgroupOk and gqaQueryCount < tunables.COOPERATIVE_QUERY_THRESHOLD and gqaCoopWorkgroupBytes <= device.limits.maxComputeWorkgroupStorageSize", + "gqaScalarDispatchOk": "gqaCoop or scalarWorkgroupOk", + "gqaDispatchUnits": "gqaQueryCount if gqaCoop else ceilDiv(gqaQueryCount, scalarWorkgroupSize)", + "gqaDispatchFits": "dim(shapes.queryT, 0) <= device.limits.maxComputeWorkgroupsPerDimension and attrs.num_heads <= device.limits.maxComputeWorkgroupsPerDimension", + "sharedKvFloatOk": "sharedKvCacheOk and f16OrF32", + "newKvPastOk": "cachedBaseOk and unquantizedCacheAttrs and not windowCacheRequested and copyWorkgroupOk and present.pastKeyT and present.pastValueT and dim(shapes.keyT, 1) > 0 and f16OrF32 and dim(shapes.queryT, 2) % attrs.num_heads == 0 and dim(shapes.pastKeyT, 1) == attrs.kv_num_heads and dim(shapes.pastKeyT, 3) == gqaHeadDim and dim(shapes.presentKeyT, 2) == dim(shapes.pastKeyT, 2) + dim(shapes.keyT, 1)", + "windowCapacity": "dim(shapes.pastKeyT, 2) if present.pastKeyT else 0", + "windowFitsCapacity": "attrs.local_window_size <= windowCapacity if attrs.local_window_size > 0 else true", + "windowShiftOk": "cachedBaseOk and unquantizedCacheAttrs and not present.cosCacheT and not rotaryRequested and windowCacheRequested and copyWorkgroupOk and present.pastKeyT and present.pastValueT and dim(shapes.keyT, 1) > 0 and f16OrF32 and dim(shapes.queryT, 2) % attrs.num_heads == 0 and dim(shapes.pastKeyT, 1) == attrs.kv_num_heads and dim(shapes.pastKeyT, 3) == gqaHeadDim and dim(shapes.presentKeyT, 2) == windowCapacity and dim(shapes.keyT, 1) <= windowCapacity and windowFitsCapacity", + "shareAppendOk": "cachedBaseOk and unquantizedCacheAttrs and not windowCacheRequested and copyWorkgroupOk and present.pastKeyT and present.pastValueT and dim(shapes.keyT, 1) > 0 and f16OrF32 and dim(shapes.queryT, 2) % attrs.num_heads == 0 and dim(shapes.pastKeyT, 1) == attrs.kv_num_heads and dim(shapes.pastKeyT, 3) == gqaHeadDim and dim(shapes.presentKeyT, 2) == dim(shapes.pastKeyT, 2) and dim(shapes.keyT, 1) <= dim(shapes.pastKeyT, 2)", + "decodeFlashShapeOk": "flashWorkgroupOk and dim(shapes.queryT, 1) == 1 and gqaHeadDim % 4 == 0 and gqaFlashHd and dim(shapes.presentKeyT, 2) >= tunables.CACHED_DECODE_MIN_KV_TOKENS and gqaDispatchFits", + "qkvPrefillClusterWorkgroupSize": "qkvPrefillQueryTile * qkvPrefillLanesPerQuery", + "cachedPrefillClusterWorkgroupSize": "cachedPrefillQueryTile * cachedPrefillLanesPerQuery", + "cachedPrefillClusterWorkgroupOk": "cachedPrefillClusterWorkgroupSize <= deviceWorkgroupCap", + "prefillFlashShapeOk": "cachedPrefillClusterWorkgroupOk and gqaHeadDim % (4 * cachedPrefillLanesPerQuery) == 0 and gqaFlashHd and dim(shapes.queryT, 1) >= cachedPrefillQueryTile and gqaDispatchFits and ceilDiv(dim(shapes.queryT, 1), cachedPrefillQueryTile) <= device.limits.maxComputeWorkgroupsPerDimension", + "quantScaleShapeOk": "((attrs.k_quant_type == \"PER_TENSOR\" and numel(shapes.kScaleT) == 1) or (attrs.k_quant_type == \"PER_CHANNEL\" and numel(shapes.kScaleT) == attrs.kv_num_heads * gqaHeadDim)) and ((attrs.v_quant_type == \"PER_TENSOR\" and numel(shapes.vScaleT) == 1) or (attrs.v_quant_type == \"PER_CHANNEL\" and numel(shapes.vScaleT) == attrs.kv_num_heads * gqaHeadDim))", + "quantizedNewKvOk": "cachedBaseOk and not windowCacheRequested and copyWorkgroupOk and present.kScaleT and present.vScaleT and attrs.k_quant_type == attrs.v_quant_type and attrs.k_quant_type != \"NONE\" and (attrs.kv_cache_bit_width == 4 or attrs.kv_cache_bit_width == 8) and quantScaleShapeOk and tensorDtypes.queryT == \"float32\" and tensorDtypes.keyT == \"float32\" and tensorDtypes.valueT == \"float32\" and tensorDtypes.outputT == \"float32\" and tensorDtypes.presentKeyT == (\"uint8\" if attrs.kv_cache_bit_width == 4 else \"int8\") and tensorDtypes.presentValueT == tensorDtypes.presentKeyT and dim(shapes.presentKeyT, 3) == gqaHeadDim / (8 / attrs.kv_cache_bit_width) and dim(shapes.presentValueT, 3) == dim(shapes.presentKeyT, 3) and dim(shapes.keyT, 1) > 0 and gqaHeadDim % 2 == 0", + "quantizedPromptOk": "quantizedNewKvOk and dim(shapes.keyT, 1) == dim(shapes.presentKeyT, 2)", + "quantizedDecodeOptionsOk": "plainAttentionOptions and not present.sinCacheT and not present.kNormWeightT", + "quantizedCachedDecodeOk": "not windowCacheRequested and quantizedNewKvOk and quantizedDecodeOptionsOk and attrs.kv_cache_bit_width == 8 and present.pastKeyT and present.pastValueT and dim(shapes.queryT, 1) == 1 and dim(shapes.keyT, 1) == 1 and tensorDtypes.pastKeyT == \"int8\" and tensorDtypes.pastValueT == \"int8\" and tensorDtypes.presentKeyT == \"int8\" and tensorDtypes.presentValueT == \"int8\" and dim(shapes.pastKeyT, 0) == dim(shapes.presentKeyT, 0) and dim(shapes.pastKeyT, 1) == dim(shapes.presentKeyT, 1) and dim(shapes.pastKeyT, 2) == dim(shapes.presentKeyT, 2) and dim(shapes.pastKeyT, 3) == dim(shapes.presentKeyT, 3)", + "qkvSplitScratchBytes": "dim(shapes.query, 0) * attrs.num_heads * qkvNumSplits * headDim * 4", + "qkvSplitStatsBytes": "2 * dim(shapes.query, 0) * attrs.num_heads * qkvNumSplits * 4", + "qkvSplitScratchFits": "qkvSplitScratchBytes <= device.limits.maxStorageBufferBindingSize and qkvSplitScratchBytes <= device.limits.maxBufferSize and qkvSplitStatsBytes <= device.limits.maxStorageBufferBindingSize and qkvSplitStatsBytes <= device.limits.maxBufferSize", + "cachedSplitScratchBytes": "dim(shapes.queryT, 0) * attrs.num_heads * cachedNumSplits * gqaHeadDim * 4", + "cachedSplitStatsBytes": "2 * dim(shapes.queryT, 0) * attrs.num_heads * cachedNumSplits * 4", + "cachedSplitScratchFits": "cachedSplitScratchBytes <= device.limits.maxStorageBufferBindingSize and cachedSplitScratchBytes <= device.limits.maxBufferSize and cachedSplitStatsBytes <= device.limits.maxStorageBufferBindingSize and cachedSplitStatsBytes <= device.limits.maxBufferSize", + "qPrepScratchBytes": "numel(shapes.queryT) * 4", + "qPrepScratchFits": "qPrepScratchBytes <= device.limits.maxStorageBufferBindingSize and qPrepScratchBytes <= device.limits.maxBufferSize", + "qkvTiledWorkgroupBytes": "dim(shapes.value, 2) / attrs.kv_num_heads * 32 * 4", + "qkvTiledStorageOk": "qkvTiledWorkgroupBytes <= device.limits.maxComputeWorkgroupStorageSize", + "qkvClusterTileK": "tunables.CLUSTER_TILE_K_F32 if tensorDtypes.query == \"float32\" and headDim <= tunables.QKV_SPLIT_TILE_K else tunables.CLUSTER_TILE_K_F16", + "qkvClusterWorkgroupBytes": "qkvClusterTileK * headDim * (8 if tensorDtypes.query == \"float32\" else 4)", + "qkvClusterStorageOk": "qkvClusterWorkgroupBytes <= device.limits.maxComputeWorkgroupStorageSize", + "qkvF32ClusterRegisterGeometry": "tensorDtypes.query != \"float32\" or (headDim % 4 == 0 and headDim / (4 * qkvPrefillLanesPerQuery) <= tunables.CLUSTER_MAX_SLICE)", + "cachedF32ClusterRegisterGeometry": "tensorDtypes.queryT != \"float32\" or (gqaHeadDim % 4 == 0 and gqaHeadDim / (4 * cachedPrefillLanesPerQuery) <= tunables.CLUSTER_MAX_SLICE)", + "qkvNoSgReductionBytesPerKey": "qkvPrefillClusterWorkgroupSize * 4", + "cachedNoSgReductionBytesPerKey": "cachedPrefillClusterWorkgroupSize * 4", + "qkvNoSgTileBytes": "headDim * (8 if tensorDtypes.query == \"float32\" else 4) + qkvNoSgReductionBytesPerKey", + "qkvNoSgTileK": "min(tunables.NO_SG_TILE_K_MAX, max(1, floor(device.limits.maxComputeWorkgroupStorageSize / qkvNoSgTileBytes)))", + "qkvNoSgWorkgroupBytes": "qkvNoSgTileK * qkvNoSgTileBytes", + "qkvNoSgClusterStorageOk": "qkvNoSgWorkgroupBytes <= device.limits.maxComputeWorkgroupStorageSize", + "cachedClusterTileK": "tunables.CLUSTER_TILE_K_F32 if tensorDtypes.queryT == \"float32\" and gqaHeadDim <= tunables.QKV_SPLIT_TILE_K else tunables.CLUSTER_TILE_K_F16", + "cachedClusterWorkgroupBytes": "cachedClusterTileK * gqaHeadDim * (8 if tensorDtypes.queryT == \"float32\" else 4)", + "cachedClusterStorageOk": "cachedClusterWorkgroupBytes <= device.limits.maxComputeWorkgroupStorageSize", + "cachedMaskClusterWorkgroupBytes": "cachedClusterWorkgroupBytes + cachedPrefillQueryTile * cachedClusterTileK * 4", + "cachedMaskClusterStorageOk": "cachedMaskClusterWorkgroupBytes <= device.limits.maxComputeWorkgroupStorageSize", + "cachedNoSgTileBytes": "gqaHeadDim * (8 if tensorDtypes.queryT == \"float32\" else 4) + cachedNoSgReductionBytesPerKey", + "cachedNoSgTileK": "min(tunables.NO_SG_TILE_K_MAX, max(1, floor(device.limits.maxComputeWorkgroupStorageSize / cachedNoSgTileBytes)))", + "cachedNoSgWorkgroupBytes": "cachedNoSgTileK * cachedNoSgTileBytes", + "cachedNoSgClusterStorageOk": "cachedNoSgWorkgroupBytes <= device.limits.maxComputeWorkgroupStorageSize" + }, + "bindingSets": { + "splitAttention": [ + { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputVec4" }, + { "name": "key", "arg": "presentKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputVec4" }, + { + "name": "value", + "arg": "presentValueT", + "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.presentKeyT, 2)" }, + { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, + { + "name": "windowSize", + "type": "u32", + "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" + } + ] + } + } + ], + "splitSeqlensAttention": [ + { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputVec4" }, + { "name": "key", "arg": "presentKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputVec4" }, + { + "name": "value", + "arg": "presentValueT", + "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": "seqlens_k", "arg": "seqlensKT", "buffer": { "type": "read-only-storage" }, "elementType": "i32" }, + { + "name": "params", + "semantic": "kernel.params", + "buffer": { "type": "uniform" }, + "struct": { + "name": "Params", + "fields": [ + { "name": "kvSeq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" }, + { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, + { + "name": "windowSize", + "type": "u32", + "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" + } + ] + } + } + ], + "splitMerge": [ + { + "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", "buffer": { "type": "storage" }, "elementType": "$inputVec4" }, + { + "name": "params", + "semantic": "kernel.params", + "buffer": { "type": "uniform" }, + "struct": { + "name": "Params", + "fields": [{ "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }] + } + } + ], + "splitMergeNoParams": [ + { + "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", "buffer": { "type": "storage" }, "elementType": "$inputVec4" } + ], + "flashPrefillPastSoftcap": [ + { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputVec4" }, + { "name": "key", "arg": "presentKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputVec4" }, + { + "name": "value", + "arg": "presentValueT", + "buffer": { "type": "read-only-storage" }, + "elementType": "$inputVec4" + }, + { "name": "output", "arg": "outputT", "buffer": { "type": "storage" }, "elementType": "$inputVec4" }, + { + "name": "params", + "semantic": "kernel.params", + "buffer": { "type": "uniform" }, + "struct": { + "name": "Params", + "fields": [ + { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, + { "name": "kvSeq", "type": "u32", "value": "dim(shapes.pastKeyT, 2)" }, + { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, + { "name": "isCausal", "type": "u32", "value": 1 }, + { + "name": "windowSize", + "type": "u32", + "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" + }, + { "name": "softcap", "type": "f32", "value": "attrs.softcap" } + ] + } + } + ], + "flashPrefillPastHeadSink": [ + { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputVec4" }, + { "name": "key", "arg": "presentKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputVec4" }, + { + "name": "value", + "arg": "presentValueT", + "buffer": { "type": "read-only-storage" }, + "elementType": "$inputVec4" + }, + { + "name": "head_sink", + "arg": "headSinkT", + "buffer": { "type": "read-only-storage" }, + "elementType": "$inputScalar" + }, + { "name": "output", "arg": "outputT", "buffer": { "type": "storage" }, "elementType": "$inputVec4" }, + { + "name": "params", + "semantic": "kernel.params", + "buffer": { "type": "uniform" }, + "struct": { + "name": "Params", + "fields": [ + { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, + { "name": "kvSeq", "type": "u32", "value": "dim(shapes.pastKeyT, 2)" }, + { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, + { "name": "isCausal", "type": "u32", "value": 1 }, + { + "name": "windowSize", + "type": "u32", + "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" + } + ] + } + } + ], + "flashPrefillBiasHeadSink": [ + { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputVec4" }, + { "name": "key", "arg": "presentKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputVec4" }, + { + "name": "value", + "arg": "presentValueT", + "buffer": { "type": "read-only-storage" }, + "elementType": "$inputVec4" + }, + { + "name": "attn_mask", + "arg": "attentionBiasT", + "buffer": { "type": "read-only-storage" }, + "elementType": "$inputScalar" + }, + { + "name": "head_sink", + "arg": "headSinkT", + "buffer": { "type": "read-only-storage" }, + "elementType": "$inputScalar" + }, + { "name": "output", "arg": "outputT", "buffer": { "type": "storage" }, "elementType": "$inputVec4" }, + { + "name": "params", + "semantic": "kernel.params", + "buffer": { "type": "uniform" }, + "struct": { + "name": "Params", + "fields": [ + { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, + { "name": "kvSeq", "type": "u32", "value": "dim(shapes.pastKeyT, 2)" }, + { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, + { "name": "isCausal", "type": "u32", "value": 1 }, + { + "name": "windowSize", + "type": "u32", + "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" + }, + { + "name": "maskBatchStride", + "type": "u32", + "value": "0 if dim(shapes.attentionBiasT, 0) == 1 else dim(shapes.attentionBiasT, 1) * dim(shapes.attentionBiasT, 2) * dim(shapes.attentionBiasT, 3)" + }, + { + "name": "maskHeadStride", + "type": "u32", + "value": "0 if dim(shapes.attentionBiasT, 1) == 1 else dim(shapes.attentionBiasT, 2) * dim(shapes.attentionBiasT, 3)" + }, + { "name": "maskSeqStride", "type": "u32", "value": "dim(shapes.attentionBiasT, 3)" } + ] + } + } + ], + "flashPrefillPast": [ + { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputVec4" }, + { "name": "key", "arg": "presentKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputVec4" }, + { + "name": "value", + "arg": "presentValueT", + "buffer": { "type": "read-only-storage" }, + "elementType": "$inputVec4" + }, + { "name": "output", "arg": "outputT", "buffer": { "type": "storage" }, "elementType": "$inputVec4" }, + { + "name": "params", + "semantic": "kernel.params", + "buffer": { "type": "uniform" }, + "struct": { + "name": "Params", + "fields": [ + { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, + { "name": "kvSeq", "type": "u32", "value": "dim(shapes.pastKeyT, 2)" }, + { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, + { "name": "isCausal", "type": "u32", "value": 1 }, + { + "name": "windowSize", + "type": "u32", + "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" + } + ] + } + } + ], + "flashPrefillSeqlens": [ + { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputVec4" }, + { "name": "key", "arg": "presentKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputVec4" }, + { + "name": "value", + "arg": "presentValueT", + "buffer": { "type": "read-only-storage" }, + "elementType": "$inputVec4" + }, + { "name": "output", "arg": "outputT", "buffer": { "type": "storage" }, "elementType": "$inputVec4" }, + { "name": "seqlens_k", "arg": "seqlensKT", "buffer": { "type": "read-only-storage" }, "elementType": "i32" }, + { + "name": "params", + "semantic": "kernel.params", + "buffer": { "type": "uniform" }, + "struct": { + "name": "Params", + "fields": [ + { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, + { "name": "kvSeq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" }, + { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, + { "name": "isCausal", "type": "u32", "value": 1 }, + { + "name": "windowSize", + "type": "u32", + "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" + } + ] + } + } + ], + "flashPrefillNew": [ + { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputVec4" }, + { "name": "key", "arg": "presentKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputVec4" }, + { + "name": "value", + "arg": "presentValueT", + "buffer": { "type": "read-only-storage" }, + "elementType": "$inputVec4" + }, + { "name": "output", "arg": "outputT", "buffer": { "type": "storage" }, "elementType": "$inputVec4" }, + { + "name": "params", + "semantic": "kernel.params", + "buffer": { "type": "uniform" }, + "struct": { + "name": "Params", + "fields": [ + { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, + { "name": "kvSeq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" }, + { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, + { "name": "isCausal", "type": "u32", "value": 1 }, + { + "name": "windowSize", + "type": "u32", + "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" + } + ] + } + } + ], + "rotaryFusedFlashAttention": [ + { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputVec4" }, + { "name": "key", "arg": "presentKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputVec4" }, + { + "name": "value", + "arg": "presentValueT", + "buffer": { "type": "read-only-storage" }, + "elementType": "$inputVec4" + }, + { + "name": "cos_cache", + "arg": "cosCacheT", + "buffer": { "type": "read-only-storage" }, + "elementType": "$cosScalar" + }, + { + "name": "sin_cache", + "arg": "sinCacheT", + "buffer": { "type": "read-only-storage" }, + "elementType": "$cosScalar" + }, + { "name": "output", "arg": "outputT", "buffer": { "type": "storage" }, "elementType": "$inputVec4" }, + { + "name": "params", + "semantic": "kernel.params", + "buffer": { "type": "uniform" }, + "struct": { + "name": "Params", + "fields": [ + { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, + { "name": "kvSeq", "type": "u32", "value": "dim(shapes.pastKeyT, 2)" }, + { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, + { "name": "isCausal", "type": "u32", "value": 1 }, + { + "name": "windowSize", + "type": "u32", + "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" + } + ] + } + } + ], + "rotaryFlashAttention": [ + { "name": "query", "semantic": "qPrep", "buffer": { "type": "read-only-storage" }, "elementType": "vec4" }, + { "name": "key", "arg": "presentKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputVec4" }, + { + "name": "value", + "arg": "presentValueT", + "buffer": { "type": "read-only-storage" }, + "elementType": "$inputVec4" + }, + { "name": "output", "arg": "outputT", "buffer": { "type": "storage" }, "elementType": "$inputVec4" }, + { + "name": "params", + "semantic": "kernel.params", + "buffer": { "type": "uniform" }, + "struct": { + "name": "Params", + "fields": [ + { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, + { "name": "kvSeq", "type": "u32", "value": "dim(shapes.pastKeyT, 2)" }, + { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, + { "name": "isCausal", "type": "u32", "value": 1 }, + { + "name": "windowSize", + "type": "u32", + "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" + } + ] + } + } + ], + "qnormRotaryQprep": [ + { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, + { + "name": "cos_cache", + "arg": "cosCacheT", + "buffer": { "type": "read-only-storage" }, + "elementType": "$cosScalar" + }, + { + "name": "sin_cache", + "arg": "sinCacheT", + "buffer": { "type": "read-only-storage" }, + "elementType": "$cosScalar" + }, + { + "name": "q_norm_weight", + "arg": "qNormWeightT", + "buffer": { "type": "read-only-storage" }, + "elementType": "$qNormScalar" + }, + { "name": "qout", "semantic": "qPrep", "buffer": { "type": "storage" }, "elementType": "vec4" }, + { + "name": "params", + "semantic": "kernel.params", + "buffer": { "type": "uniform" }, + "struct": { + "name": "Params", + "fields": [ + { "name": "batch", "type": "u32", "value": "dim(shapes.queryT, 0)" }, + { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, + { "name": "totalSeq", "type": "u32", "value": "dim(shapes.pastKeyT, 2)" } + ] + } + } + ], + "biasFlashAttention": [ + { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputVec4" }, + { "name": "key", "arg": "presentKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputVec4" }, + { + "name": "value", + "arg": "presentValueT", + "buffer": { "type": "read-only-storage" }, + "elementType": "$inputVec4" + }, + { + "name": "attn_mask", + "arg": "attentionBiasT", + "buffer": { "type": "read-only-storage" }, + "elementType": "$inputScalar" + }, + { "name": "output", "arg": "outputT", "buffer": { "type": "storage" }, "elementType": "$inputVec4" }, + { + "name": "params", + "semantic": "kernel.params", + "buffer": { "type": "uniform" }, + "struct": { + "name": "Params", + "fields": [ + { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, + { "name": "kvSeq", "type": "u32", "value": "dim(shapes.pastKeyT, 2)" }, + { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, + { "name": "isCausal", "type": "u32", "value": 1 }, + { + "name": "windowSize", + "type": "u32", + "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" + }, + { + "name": "maskBatchStride", + "type": "u32", + "value": "0 if dim(shapes.attentionBiasT, 0) == 1 else dim(shapes.attentionBiasT, 1) * dim(shapes.attentionBiasT, 2) * dim(shapes.attentionBiasT, 3)" + }, + { + "name": "maskHeadStride", + "type": "u32", + "value": "0 if dim(shapes.attentionBiasT, 1) == 1 else dim(shapes.attentionBiasT, 2) * dim(shapes.attentionBiasT, 3)" + }, + { "name": "maskSeqStride", "type": "u32", "value": "dim(shapes.attentionBiasT, 3)" } + ] + } + } + ], + "qkvSplitAttention": [ + { + "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 attrs.scale else 0" }, + { + "name": "windowSize", + "type": "u32", + "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" + } + ] + } + } + ], + "qkvSplitMerge": [ + { + "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 attrs.scale else 0" } + ] + } + } + ], + "qkvSplitMergeNoParams": [ + { + "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" + } + ], + "qkvPresent": [ + { + "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": "present_key", + "arg": "presentKeyT", + "semantic": "present_key", + "buffer": { "type": "storage" }, + "elementType": "$presentElement" + }, + { + "name": "present_value", + "arg": "presentValueT", + "semantic": "present_value", + "buffer": { "type": "storage" }, + "elementType": "$presentElement" + }, + { + "name": "params", + "semantic": "kernel.params", + "buffer": { "type": "uniform" }, + "struct": { + "name": "Params", + "fields": [ + { "name": "batchSize", "type": "u32", "value": "dim(shapes.query, 0)" }, + { "name": "kvSeq", "type": "u32", "value": "dim(shapes.key, 1)" } + ] + } + } + ], + "qkvAttention": [ + { + "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 attrs.scale else 0" }, + { + "name": "windowSize", + "type": "u32", + "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" + } + ] + } + } + ], + "qkvAttentionNoWindow": [ + { + "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 attrs.scale else 0" } + ] + } + } + ], + "presentCopy": [ + { "name": "src_k", "arg": "pastKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, + { "name": "src_v", "arg": "pastValueT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, + { "name": "present_key", "arg": "presentKeyT", "buffer": { "type": "storage" }, "elementType": "$inputScalar" }, + { + "name": "present_value", + "arg": "presentValueT", + "buffer": { "type": "storage" }, + "elementType": "$inputScalar" + }, + { + "name": "params", + "semantic": "kernel.params", + "buffer": { "type": "uniform" }, + "struct": { + "name": "Params", + "fields": [ + { + "name": "count", + "type": "u32", + "value": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)" + } + ] + } + } + ], + "presentMergeShare": [ + { "name": "past_k", "arg": "pastKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, + { + "name": "past_v", + "arg": "pastValueT", + "buffer": { "type": "read-only-storage" }, + "elementType": "$inputScalar" + }, + { "name": "new_k", "arg": "keyT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, + { "name": "new_v", "arg": "valueT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, + { "name": "present_key", "arg": "presentKeyT", "buffer": { "type": "storage" }, "elementType": "$inputScalar" }, + { + "name": "present_value", + "arg": "presentValueT", + "buffer": { "type": "storage" }, + "elementType": "$inputScalar" + }, + { "name": "seqlens_k", "arg": "seqlensKT", "buffer": { "type": "read-only-storage" }, "elementType": "i32" }, + { + "name": "params", + "semantic": "kernel.params", + "buffer": { "type": "uniform" }, + "struct": { + "name": "Params", + "fields": [ + { + "name": "count", + "type": "u32", + "value": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)" + }, + { "name": "seq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" }, + { "name": "keySeq", "type": "u32", "value": "dim(shapes.keyT, 1)" } + ] + } + } + ], + "presentMergeShareRetain": [ + { "name": "past_k", "arg": "pastKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, + { + "name": "past_v", + "arg": "pastValueT", + "buffer": { "type": "read-only-storage" }, + "elementType": "$inputScalar" + }, + { "name": "present_key", "arg": "presentKeyT", "buffer": { "type": "storage" }, "elementType": "$inputScalar" }, + { + "name": "present_value", + "arg": "presentValueT", + "buffer": { "type": "storage" }, + "elementType": "$inputScalar" + }, + { "name": "seqlens_k", "arg": "seqlensKT", "buffer": { "type": "read-only-storage" }, "elementType": "i32" }, + { + "name": "params", + "semantic": "kernel.params", + "buffer": { "type": "uniform" }, + "struct": { + "name": "Params", + "fields": [ + { + "name": "count", + "type": "u32", + "value": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)" + }, + { "name": "seq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" }, + { "name": "keySeq", "type": "u32", "value": "dim(shapes.keyT, 1)" } + ] + } + } + ], + "presentMergeShareAppend": [ + { "name": "new_k", "arg": "keyT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, + { "name": "new_v", "arg": "valueT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, + { "name": "present_key", "arg": "presentKeyT", "buffer": { "type": "storage" }, "elementType": "$inputScalar" }, + { + "name": "present_value", + "arg": "presentValueT", + "buffer": { "type": "storage" }, + "elementType": "$inputScalar" + }, + { "name": "seqlens_k", "arg": "seqlensKT", "buffer": { "type": "read-only-storage" }, "elementType": "i32" }, + { + "name": "params", + "semantic": "kernel.params", + "buffer": { "type": "uniform" }, + "struct": { + "name": "Params", + "fields": [ + { + "name": "count", + "type": "u32", + "value": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)" + }, + { "name": "seq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" }, + { "name": "keySeq", "type": "u32", "value": "dim(shapes.keyT, 1)" } + ] + } + } + ], + "presentMergeShareRotaryAppend": [ + { "name": "new_k", "arg": "keyT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, + { "name": "new_v", "arg": "valueT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, + { "name": "cos_cache", "arg": "cosCacheT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, + { "name": "sin_cache", "arg": "sinCacheT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, + { "name": "present_key", "arg": "presentKeyT", "buffer": { "type": "storage" }, "elementType": "$inputScalar" }, + { + "name": "present_value", + "arg": "presentValueT", + "buffer": { "type": "storage" }, + "elementType": "$inputScalar" + }, + { "name": "seqlens_k", "arg": "seqlensKT", "buffer": { "type": "read-only-storage" }, "elementType": "i32" }, + { + "name": "params", + "semantic": "kernel.params", + "buffer": { "type": "uniform" }, + "struct": { + "name": "Params", + "fields": [ + { + "name": "count", + "type": "u32", + "value": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)" + }, + { "name": "seq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" }, + { "name": "keySeq", "type": "u32", "value": "dim(shapes.keyT, 1)" } + ] + } + } + ], + "mergedSeqlensAttention": [ + { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, + { + "name": "kcache", + "arg": "presentKeyT", + "buffer": { "type": "read-only-storage" }, + "elementType": "$inputScalar" + }, + { + "name": "vcache", + "arg": "presentValueT", + "buffer": { "type": "read-only-storage" }, + "elementType": "$inputScalar" + }, + { "name": "output", "arg": "outputT", "buffer": { "type": "storage" }, "elementType": "$inputScalar" }, + { "name": "seqlens_k", "arg": "seqlensKT", "buffer": { "type": "read-only-storage" }, "elementType": "i32" }, + { + "name": "params", + "semantic": "kernel.params", + "buffer": { "type": "uniform" }, + "struct": { + "name": "Params", + "fields": [ + { "name": "batch", "type": "u32", "value": "dim(shapes.queryT, 0)" }, + { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, + { "name": "totalSeq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" }, + { + "name": "windowSize", + "type": "u32", + "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" + }, + { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, + { "name": "softcap", "type": "f32", "value": "attrs.softcap" }, + { "name": "smoothSoftmax", "type": "u32", "value": "1 if attrs.smooth_softmax == 1 else 0" } + ] + } + } + ], + "mergedSeqlensHeadSinkAttention": [ + { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, + { + "name": "kcache", + "arg": "presentKeyT", + "buffer": { "type": "read-only-storage" }, + "elementType": "$inputScalar" + }, + { + "name": "vcache", + "arg": "presentValueT", + "buffer": { "type": "read-only-storage" }, + "elementType": "$inputScalar" + }, + { + "name": "head_sink", + "arg": "headSinkT", + "buffer": { "type": "read-only-storage" }, + "elementType": "$inputScalar" + }, + { "name": "output", "arg": "outputT", "buffer": { "type": "storage" }, "elementType": "$inputScalar" }, + { "name": "seqlens_k", "arg": "seqlensKT", "buffer": { "type": "read-only-storage" }, "elementType": "i32" }, + { + "name": "params", + "semantic": "kernel.params", + "buffer": { "type": "uniform" }, + "struct": { + "name": "Params", + "fields": [ + { "name": "batch", "type": "u32", "value": "dim(shapes.queryT, 0)" }, + { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, + { "name": "totalSeq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" }, + { + "name": "windowSize", + "type": "u32", + "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" + }, + { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, + { "name": "softcap", "type": "f32", "value": "attrs.softcap" } + ] + } + } + ], + "mergedSeqlensRotaryAttention": [ + { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, + { + "name": "kcache", + "arg": "presentKeyT", + "buffer": { "type": "read-only-storage" }, + "elementType": "$inputScalar" + }, + { + "name": "vcache", + "arg": "presentValueT", + "buffer": { "type": "read-only-storage" }, + "elementType": "$inputScalar" + }, + { "name": "cos_cache", "arg": "cosCacheT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, + { "name": "sin_cache", "arg": "sinCacheT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, + { "name": "output", "arg": "outputT", "buffer": { "type": "storage" }, "elementType": "$inputScalar" }, + { "name": "seqlens_k", "arg": "seqlensKT", "buffer": { "type": "read-only-storage" }, "elementType": "i32" }, + { + "name": "params", + "semantic": "kernel.params", + "buffer": { "type": "uniform" }, + "struct": { + "name": "Params", + "fields": [ + { "name": "batch", "type": "u32", "value": "dim(shapes.queryT, 0)" }, + { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, + { "name": "totalSeq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" }, + { + "name": "windowSize", + "type": "u32", + "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" + }, + { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, + { "name": "softcap", "type": "f32", "value": "attrs.softcap" }, + { "name": "smoothSoftmax", "type": "u32", "value": "1 if attrs.smooth_softmax == 1 else 0" } + ] + } + } + ], + "presentMerge": [ + { "name": "past_k", "arg": "pastKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, + { + "name": "past_v", + "arg": "pastValueT", + "buffer": { "type": "read-only-storage" }, + "elementType": "$inputScalar" + }, + { "name": "new_k", "arg": "keyT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, + { "name": "new_v", "arg": "valueT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, + { "name": "present_key", "arg": "presentKeyT", "buffer": { "type": "storage" }, "elementType": "$inputScalar" }, + { + "name": "present_value", + "arg": "presentValueT", + "buffer": { "type": "storage" }, + "elementType": "$inputScalar" + }, + { + "name": "params", + "semantic": "kernel.params", + "buffer": { "type": "uniform" }, + "struct": { + "name": "Params", + "fields": [ + { + "name": "count", + "type": "u32", + "value": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)" + }, + { "name": "seq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" }, + { "name": "pastSeq", "type": "u32", "value": "dim(shapes.pastKeyT, 2)" }, + { "name": "keySeq", "type": "u32", "value": "dim(shapes.keyT, 1)" } + ] + } + } + ], + "qkvTiledAttention": [ + { + "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.outputT)" }, + { "name": "qHeads", "type": "u32", "value": "attrs.num_heads" }, + { "name": "kvHeads", "type": "u32", "value": "attrs.kv_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.kv_num_heads" }, + { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, + { "name": "softcap", "type": "f32", "value": "0" }, + { "name": "isCausal", "type": "u32", "value": "0" }, + { "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)" } + ] + } + } + ], + "cachedAttention": [ + { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, + { "name": "kcache", "arg": "pastKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, + { + "name": "vcache", + "arg": "pastValueT", + "buffer": { "type": "read-only-storage" }, + "elementType": "$inputScalar" + }, + { "name": "output", "arg": "outputT", "buffer": { "type": "storage" }, "elementType": "$inputScalar" }, + { "name": "seqlens_k", "arg": "seqlensKT", "buffer": { "type": "read-only-storage" }, "elementType": "i32" }, + { + "name": "params", + "semantic": "kernel.params", + "buffer": { "type": "uniform" }, + "struct": { + "name": "Params", + "fields": [ + { "name": "batch", "type": "u32", "value": "dim(shapes.queryT, 0)" }, + { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, + { "name": "totalSeq", "type": "u32", "value": "dim(shapes.pastKeyT, 2)" }, + { + "name": "windowSize", + "type": "u32", + "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" + }, + { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, + { "name": "softcap", "type": "f32", "value": "attrs.softcap" }, + { "name": "smoothSoftmax", "type": "u32", "value": "1 if attrs.smooth_softmax == 1 else 0" } + ] + } + } + ], + "cachedRotaryAttention": [ + { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, + { "name": "kcache", "arg": "pastKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, + { + "name": "vcache", + "arg": "pastValueT", + "buffer": { "type": "read-only-storage" }, + "elementType": "$inputScalar" + }, + { + "name": "cos_cache", + "arg": "cosCacheT", + "buffer": { "type": "read-only-storage" }, + "elementType": "$cosScalar" + }, + { + "name": "sin_cache", + "arg": "sinCacheT", + "buffer": { "type": "read-only-storage" }, + "elementType": "$cosScalar" + }, + { "name": "output", "arg": "outputT", "buffer": { "type": "storage" }, "elementType": "$inputScalar" }, + { "name": "seqlens_k", "arg": "seqlensKT", "buffer": { "type": "read-only-storage" }, "elementType": "i32" }, + { + "name": "params", + "semantic": "kernel.params", + "buffer": { "type": "uniform" }, + "struct": { + "name": "Params", + "fields": [ + { "name": "batch", "type": "u32", "value": "dim(shapes.queryT, 0)" }, + { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, + { "name": "totalSeq", "type": "u32", "value": "dim(shapes.pastKeyT, 2)" }, + { + "name": "windowSize", + "type": "u32", + "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" + }, + { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, + { "name": "softcap", "type": "f32", "value": "attrs.softcap" }, + { "name": "smoothSoftmax", "type": "u32", "value": "1 if attrs.smooth_softmax == 1 else 0" } + ] + } + } + ], + "cachedQnormRotaryAttention": [ + { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, + { "name": "kcache", "arg": "pastKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, + { + "name": "vcache", + "arg": "pastValueT", + "buffer": { "type": "read-only-storage" }, + "elementType": "$inputScalar" + }, + { + "name": "cos_cache", + "arg": "cosCacheT", + "buffer": { "type": "read-only-storage" }, + "elementType": "$cosScalar" + }, + { + "name": "sin_cache", + "arg": "sinCacheT", + "buffer": { "type": "read-only-storage" }, + "elementType": "$cosScalar" + }, + { + "name": "q_norm_weight", + "arg": "qNormWeightT", + "buffer": { "type": "read-only-storage" }, + "elementType": "$qNormScalar" + }, + { "name": "output", "arg": "outputT", "buffer": { "type": "storage" }, "elementType": "$inputScalar" }, + { "name": "seqlens_k", "arg": "seqlensKT", "buffer": { "type": "read-only-storage" }, "elementType": "i32" }, + { + "name": "params", + "semantic": "kernel.params", + "buffer": { "type": "uniform" }, + "struct": { + "name": "Params", + "fields": [ + { "name": "batch", "type": "u32", "value": "dim(shapes.queryT, 0)" }, + { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, + { "name": "totalSeq", "type": "u32", "value": "dim(shapes.pastKeyT, 2)" }, + { + "name": "windowSize", + "type": "u32", + "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" + }, + { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, + { "name": "softcap", "type": "f32", "value": "attrs.softcap" }, + { "name": "smoothSoftmax", "type": "u32", "value": "1 if attrs.smooth_softmax == 1 else 0" } + ] + } + } + ], + "newQnormRotaryPresent": [ + { "name": "src_k", "arg": "keyT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, + { "name": "src_v", "arg": "valueT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, + { "name": "cos_cache", "arg": "cosCacheT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, + { "name": "sin_cache", "arg": "sinCacheT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, + { + "name": "k_norm_weight", + "arg": "kNormWeightT", + "buffer": { "type": "read-only-storage" }, + "elementType": "f32" + }, + { "name": "present_key", "arg": "presentKeyT", "buffer": { "type": "storage" }, "elementType": "f32" }, + { "name": "present_value", "arg": "presentValueT", "buffer": { "type": "storage" }, "elementType": "f32" }, + { + "name": "params", + "semantic": "kernel.params", + "buffer": { "type": "uniform" }, + "struct": { + "name": "Params", + "fields": [ + { + "name": "count", + "type": "u32", + "value": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)" + }, + { "name": "seq", "type": "u32", "value": "dim(shapes.keyT, 1)" }, + { "name": "pastSeq", "type": "u32", "value": 0 } + ] + } + } + ], + "newQnormRotaryAttention": [ + { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, + { "name": "kcache", "arg": "presentKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, + { "name": "vcache", "arg": "presentValueT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, + { "name": "cos_cache", "arg": "cosCacheT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, + { "name": "sin_cache", "arg": "sinCacheT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, + { + "name": "q_norm_weight", + "arg": "qNormWeightT", + "buffer": { "type": "read-only-storage" }, + "elementType": "f32" + }, + { "name": "output", "arg": "outputT", "buffer": { "type": "storage" }, "elementType": "f32" }, + { + "name": "params", + "semantic": "kernel.params", + "buffer": { "type": "uniform" }, + "struct": { + "name": "Params", + "fields": [ + { "name": "batch", "type": "u32", "value": "dim(shapes.queryT, 0)" }, + { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, + { "name": "totalSeq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" }, + { + "name": "windowSize", + "type": "u32", + "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" + }, + { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, + { "name": "softcap", "type": "f32", "value": "attrs.softcap" }, + { "name": "smoothSoftmax", "type": "u32", "value": "1 if attrs.smooth_softmax == 1 else 0" } + ] + } + } + ], + "mergedAttention": [ + { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, + { + "name": "kcache", + "arg": "presentKeyT", + "buffer": { "type": "read-only-storage" }, + "elementType": "$inputScalar" + }, + { + "name": "vcache", + "arg": "presentValueT", + "buffer": { "type": "read-only-storage" }, + "elementType": "$inputScalar" + }, + { "name": "output", "arg": "outputT", "buffer": { "type": "storage" }, "elementType": "$inputScalar" }, + { + "name": "params", + "semantic": "kernel.params", + "buffer": { "type": "uniform" }, + "struct": { + "name": "Params", + "fields": [ + { "name": "batch", "type": "u32", "value": "dim(shapes.queryT, 0)" }, + { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, + { "name": "totalSeq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" }, + { + "name": "windowSize", + "type": "u32", + "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" + }, + { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, + { "name": "softcap", "type": "f32", "value": "attrs.softcap" }, + { "name": "smoothSoftmax", "type": "u32", "value": "1 if attrs.smooth_softmax == 1 else 0" } + ] + } + } + ], + "quantPresent": [ + { "name": "src_k", "arg": "keyT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, + { "name": "src_v", "arg": "valueT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, + { "name": "k_scale", "arg": "kScaleT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, + { "name": "v_scale", "arg": "vScaleT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, + { + "name": "present_key", + "arg": "presentKeyT", + "buffer": { "type": "storage" }, + "elementType": "$quantCacheElement" + }, + { + "name": "present_value", + "arg": "presentValueT", + "buffer": { "type": "storage" }, + "elementType": "$quantCacheElement" + }, + { + "name": "params", + "semantic": "kernel.params", + "buffer": { "type": "uniform" }, + "struct": { + "name": "Params", + "fields": [ + { + "name": "count", + "type": "u32", + "value": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)" + }, + { "name": "seq", "type": "u32", "value": "dim(shapes.keyT, 1)" }, + { "name": "perChannel", "type": "u32", "value": "1 if dim(shapes.kScaleT, 0) > 1 else 0" } + ] + } + } + ], + "quantInt8DecodePresent": [ + { "name": "src_k", "arg": "keyT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, + { "name": "src_v", "arg": "valueT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, + { "name": "k_scale", "arg": "kScaleT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, + { "name": "v_scale", "arg": "vScaleT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, + { "name": "seqlens_k", "arg": "seqlensKT", "buffer": { "type": "read-only-storage" }, "elementType": "i32" }, + { "name": "present_key", "arg": "presentKeyT", "buffer": { "type": "storage" }, "elementType": "i32" }, + { "name": "present_value", "arg": "presentValueT", "buffer": { "type": "storage" }, "elementType": "i32" }, + { + "name": "params", + "semantic": "kernel.params", + "buffer": { "type": "uniform" }, + "struct": { + "name": "Params", + "fields": [ + { + "name": "count", + "type": "u32", + "value": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)" + }, + { "name": "seq", "type": "u32", "value": "dim(shapes.keyT, 1)" }, + { "name": "totalSeq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" }, + { "name": "perChannel", "type": "u32", "value": "1 if dim(shapes.kScaleT, 0) > 1 else 0" } + ] + } + } + ], + "quantAttention": [ + { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, + { + "name": "kcache", + "arg": "presentKeyT", + "buffer": { "type": "read-only-storage" }, + "elementType": "$quantCacheElement" + }, + { + "name": "vcache", + "arg": "presentValueT", + "buffer": { "type": "read-only-storage" }, + "elementType": "$quantCacheElement" + }, + { "name": "k_scale", "arg": "kScaleT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, + { "name": "v_scale", "arg": "vScaleT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, + { "name": "output", "arg": "outputT", "buffer": { "type": "storage" }, "elementType": "f32" }, + { + "name": "params", + "semantic": "kernel.params", + "buffer": { "type": "uniform" }, + "struct": { + "name": "Params", + "fields": [ + { "name": "batch", "type": "u32", "value": "dim(shapes.queryT, 0)" }, + { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, + { "name": "totalSeq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" }, + { + "name": "windowSize", + "type": "u32", + "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" + }, + { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, + { "name": "softcap", "type": "f32", "value": "attrs.softcap" }, + { "name": "smoothSoftmax", "type": "u32", "value": "1 if attrs.smooth_softmax == 1 else 0" }, + { "name": "perChannel", "type": "u32", "value": "1 if dim(shapes.kScaleT, 0) > 1 else 0" } + ] + } + } + ], + "quantInt8SplitAttention": [ + { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "vec4" }, + { "name": "key", "arg": "presentKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "vec4" }, + { "name": "value", "arg": "presentValueT", "buffer": { "type": "read-only-storage" }, "elementType": "vec4" }, + { "name": "k_scale", "arg": "kScaleT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, + { "name": "v_scale", "arg": "vScaleT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, + { "name": "seqlens_k", "arg": "seqlensKT", "buffer": { "type": "read-only-storage" }, "elementType": "i32" }, + { "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.presentKeyT, 2)" }, + { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, + { + "name": "windowSize", + "type": "u32", + "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" + }, + { "name": "perChannel", "type": "u32", "value": "1 if dim(shapes.kScaleT, 0) > 1 else 0" } + ] + } + } + ], + "cachedBiasAttention": [ + { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, + { "name": "kcache", "arg": "pastKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, + { "name": "vcache", "arg": "pastValueT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, + { + "name": "attn_bias", + "arg": "attentionBiasT", + "buffer": { "type": "read-only-storage" }, + "elementType": "$inputScalar" + }, + { "name": "output", "arg": "outputT", "buffer": { "type": "storage" }, "elementType": "f32" }, + { "name": "seqlens_k", "arg": "seqlensKT", "buffer": { "type": "read-only-storage" }, "elementType": "i32" }, + { + "name": "params", + "semantic": "kernel.params", + "buffer": { "type": "uniform" }, + "struct": { + "name": "Params", + "fields": [ + { "name": "batch", "type": "u32", "value": "dim(shapes.queryT, 0)" }, + { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, + { "name": "totalSeq", "type": "u32", "value": "dim(shapes.pastKeyT, 2)" }, + { + "name": "windowSize", + "type": "u32", + "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" + }, + { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, + { "name": "softcap", "type": "f32", "value": "attrs.softcap" }, + { "name": "smoothSoftmax", "type": "u32", "value": "1 if attrs.smooth_softmax == 1 else 0" }, + { "name": "biasBatch", "type": "u32", "value": "dim(shapes.attentionBiasT, 0)" }, + { "name": "biasHeads", "type": "u32", "value": "dim(shapes.attentionBiasT, 1)" } + ] + } + } + ], + "cachedHeadSinkAttention": [ + { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, + { "name": "kcache", "arg": "pastKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, + { "name": "vcache", "arg": "pastValueT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, + { + "name": "head_sink", + "arg": "headSinkT", + "buffer": { "type": "read-only-storage" }, + "elementType": "$inputScalar" + }, + { "name": "output", "arg": "outputT", "buffer": { "type": "storage" }, "elementType": "f32" }, + { "name": "seqlens_k", "arg": "seqlensKT", "buffer": { "type": "read-only-storage" }, "elementType": "i32" }, + { + "name": "params", + "semantic": "kernel.params", + "buffer": { "type": "uniform" }, + "struct": { + "name": "Params", + "fields": [ + { "name": "batch", "type": "u32", "value": "dim(shapes.queryT, 0)" }, + { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, + { "name": "totalSeq", "type": "u32", "value": "dim(shapes.pastKeyT, 2)" }, + { + "name": "windowSize", + "type": "u32", + "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" + }, + { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, + { "name": "softcap", "type": "f32", "value": "attrs.softcap" } + ] + } + } + ], + "cachedBiasHeadSinkAttention": [ + { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, + { "name": "kcache", "arg": "pastKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, + { "name": "vcache", "arg": "pastValueT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, + { + "name": "attn_bias", + "arg": "attentionBiasT", + "buffer": { "type": "read-only-storage" }, + "elementType": "$inputScalar" + }, + { + "name": "head_sink", + "arg": "headSinkT", + "buffer": { "type": "read-only-storage" }, + "elementType": "$inputScalar" + }, + { "name": "output", "arg": "outputT", "buffer": { "type": "storage" }, "elementType": "f32" }, + { "name": "seqlens_k", "arg": "seqlensKT", "buffer": { "type": "read-only-storage" }, "elementType": "i32" }, + { + "name": "params", + "semantic": "kernel.params", + "buffer": { "type": "uniform" }, + "struct": { + "name": "Params", + "fields": [ + { "name": "batch", "type": "u32", "value": "dim(shapes.queryT, 0)" }, + { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, + { "name": "totalSeq", "type": "u32", "value": "dim(shapes.pastKeyT, 2)" }, + { + "name": "windowSize", + "type": "u32", + "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" + }, + { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, + { "name": "softcap", "type": "f32", "value": "attrs.softcap" }, + { "name": "biasBatch", "type": "u32", "value": "dim(shapes.attentionBiasT, 0)" }, + { "name": "biasHeads", "type": "u32", "value": "dim(shapes.attentionBiasT, 1)" } + ] + } + } + ], + "flashPrefillQuant": [ + { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputVec4" }, + { + "name": "key", + "arg": "presentKeyT", + "buffer": { "type": "read-only-storage" }, + "elementType": "$quantCacheReadElement" + }, + { + "name": "value", + "arg": "presentValueT", + "buffer": { "type": "read-only-storage" }, + "elementType": "$quantCacheReadElement" + }, + { "name": "k_scale", "arg": "kScaleT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, + { "name": "v_scale", "arg": "vScaleT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, + { "name": "output", "arg": "outputT", "buffer": { "type": "storage" }, "elementType": "$inputVec4" }, + { + "name": "params", + "semantic": "kernel.params", + "buffer": { "type": "uniform" }, + "struct": { + "name": "Params", + "fields": [ + { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, + { "name": "kvSeq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" }, + { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, + { "name": "isCausal", "type": "u32", "value": 1 }, + { + "name": "windowSize", + "type": "u32", + "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" + }, + { "name": "perChannel", "type": "u32", "value": "1 if dim(shapes.kScaleT, 0) > 1 else 0" } + ] + } + } + ] + }, + "variants": [ + { + "id": "new_kv_share_append_split", + "description": "Retains the existing cache and appends new key/value rows in separate passes before portable attention. It avoids rebuilding unchanged cache positions when the past and present allocations share capacity.", + "priority": 32, + "when": ["shareAppendOk", "gqaScalarDispatchOk", "noAuxAttentionInputs", "not present.kNormWeightT", "gqaHeadDim % 2 == 0"], + "constants": { + "useSeqlens": true, + "headDim": "gqaHeadDim", + "cooperative": "gqaCoop", + "dispatchUnits": "gqaDispatchUnits", + "qHeads": "attrs.num_heads", + "kvHeads": "attrs.kv_num_heads", + "qHidden": "dim(shapes.queryT, 2)", + "packed": "gqaHeadDim", + "mode": "\"merge_share\"", + "usesF16": "tensorDtypes.queryT == \"float16\"", + "inputScalar": "gqaScalar" + }, + "passes": [ + { + "id": "present_retain", + "name": "GroupQueryAttention.MergeShareRetain", + "shader": "gqa-present.wgsl.jinja", + "bindings": "presentMergeShareRetain", + "constants": { "shareRegion": "\"retain\"" }, + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)", + "workgroupSize": "copyWorkgroupSize" + }, + "viewAlias": [ + { "input": "past_k", "output": "present_key" }, + { "input": "past_v", "output": "present_value" } + ] + }, + { + "id": "present_append", + "name": "GroupQueryAttention.MergeShareAppend", + "shader": "gqa-present.wgsl.jinja", + "bindings": "presentMergeShareAppend", + "constants": { "shareRegion": "\"append\"" }, + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)", + "workgroupSize": "copyWorkgroupSize" + } + }, + { + "id": "attention", + "name": "GroupQueryAttention.Attn", + "shader": "gqa-attention.wgsl.jinja", + "bindings": "mergedSeqlensAttention", + "dispatch": { "workgroups": "constants.dispatchUnits" } + } + ] + }, + { + "id": "new_kv_share_append_headsink_split", + "priority": 30, + "when": ["shareAppendOk", "gqaScalarDispatchOk", "gqaHeadDim % 2 == 0", "not present.cosCacheT and not rotaryRequested", "not present.qNormWeightT", "not present.kNormWeightT", "not present.attentionBiasT", "present.headSinkT"], + "constants": { + "useSeqlens": true, + "headDim": "gqaHeadDim", + "cooperative": "gqaCoop", + "dispatchUnits": "gqaDispatchUnits", + "qHeads": "attrs.num_heads", + "kvHeads": "attrs.kv_num_heads", + "qHidden": "dim(shapes.queryT, 2)", + "packed": "gqaHeadDim", + "hasHeadSink": true, + "mode": "\"merge_share\"", + "usesF16": "tensorDtypes.queryT == \"float16\"", + "inputScalar": "gqaScalar" + }, + "passes": [ + { + "id": "present_retain", + "name": "GroupQueryAttention.MergeShareRetain", + "shader": "gqa-present.wgsl.jinja", + "bindings": "presentMergeShareRetain", + "constants": { "shareRegion": "\"retain\"" }, + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)", + "workgroupSize": "copyWorkgroupSize" + }, + "viewAlias": [ + { "input": "past_k", "output": "present_key" }, + { "input": "past_v", "output": "present_value" } + ] + }, + { + "id": "present_append", + "name": "GroupQueryAttention.MergeShareAppend", + "shader": "gqa-present.wgsl.jinja", + "bindings": "presentMergeShareAppend", + "constants": { "shareRegion": "\"append\"" }, + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)", + "workgroupSize": "copyWorkgroupSize" + } + }, + { + "id": "attention", + "name": "GroupQueryAttention.Attn", + "shader": "gqa-attention.wgsl.jinja", + "bindings": "mergedSeqlensHeadSinkAttention", + "dispatch": { "workgroups": "constants.dispatchUnits" } + } + ] + }, + { + "id": "new_kv_share_append_rotary_split", + "priority": 30, + "when": ["shareAppendOk", "gqaScalarDispatchOk", "gqaHeadDim % 2 == 0", "present.cosCacheT and rotaryRequested", "not present.qNormWeightT", "not present.kNormWeightT", "not present.attentionBiasT", "not present.headSinkT", "tensorDtypes.queryT == \"float32\""], + "constants": { + "useSeqlens": true, + "headDim": "gqaHeadDim", + "cooperative": "gqaCoop", + "dispatchUnits": "gqaDispatchUnits", + "qHeads": "attrs.num_heads", + "kvHeads": "attrs.kv_num_heads", + "qHidden": "dim(shapes.queryT, 2)", + "packed": "gqaHeadDim", + "half": "gqaHeadDim / 2", + "hasRotary": true, + "mode": "\"merge_share\"", + "inputScalar": "\"f32\"", + "cosScalar": "\"f32\"" + }, + "passes": [ + { + "id": "present_retain", + "name": "GroupQueryAttention.MergeShareRetain", + "shader": "gqa-present.wgsl.jinja", + "bindings": "presentMergeShareRetain", + "constants": { "shareRegion": "\"retain\"" }, + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)", + "workgroupSize": "copyWorkgroupSize" + }, + "viewAlias": [ + { "input": "past_k", "output": "present_key" }, + { "input": "past_v", "output": "present_value" } + ] + }, + { + "id": "present_append", + "name": "GroupQueryAttention.MergeShareAppend", + "shader": "gqa-present.wgsl.jinja", + "bindings": "presentMergeShareRotaryAppend", + "constants": { "shareRegion": "\"append\"" }, + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)", + "workgroupSize": "copyWorkgroupSize" + } + }, + { + "id": "attention", + "name": "GroupQueryAttention.Attn", + "shader": "gqa-attention.wgsl.jinja", + "bindings": "mergedSeqlensRotaryAttention", + "dispatch": { "workgroups": "constants.dispatchUnits" } + } + ] + }, + { + "id": "qkv_present_tiled_nosg", + "description": "Portable tiled prefill route that computes attention online and writes the present cache separately. It is used when the flash shape is valid but no suitable subgroup route is admissible.", + "priority": 19, + "when": ["plainAttentionOptions", "flashShapeOk", "qkvTiledStorageOk", "headDim % 4 == 0", "headDim <= 128", "dim(shapes.query, 1) >= tunables.QKV_PREFILL_MIN_QUERY_TOKENS", "attrs.local_window_size <= 0"], + "constants": { + "scalar": "qkvScalar", + "inputScalar": "qkvScalar", + "outputScalar": "qkvScalar", + "presentScalar": "qkvCacheScalar", + "inputElement": "qkvScalar", + "outputElement": "qkvScalar", + "presentElement": "qkvCacheScalar", + "usesF16": "qkvUsesF16", + "blockM": 32, + "vHeadCap": "dim(shapes.value, 2) / attrs.kv_num_heads", + "kvNumHeads": "attrs.kv_num_heads", + "headDim": "dim(shapes.query, 2) / attrs.num_heads", + "qHidden": "dim(shapes.query, 2)", + "kvHidden": "dim(shapes.key, 2)", + "copyWorkgroupSize": "copyWorkgroupSize" + }, + "passes": [ + { + "id": "attention", + "name": "GroupQueryAttention.PrefillTiledNoSg", + "source": { "shader": "attention-rank4-tiled.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, + "constants": { "usesF16": "tensorDtypes.query == \"float16\"" }, + "bindings": "qkvTiledAttention", + "dispatch": { + "threads": "dim(shapes.outputT, 0) * attrs.num_heads * ceilDiv(dim(shapes.outputT, 1), 32) * 32", + "workgroupSize": "32" + } + }, + { + "id": "present", + "name": "GroupQueryAttention.PresentTiledNoSg", + "shader": "gqa-present.wgsl.jinja", + "bindings": "qkvPresent", + "dispatch": { "threads": "numel(shapes.presentKeyT)", "workgroupSize": "constants.copyWorkgroupSize" } + } + ] + }, + { + "id": "qkv_present_flash", + "priority": 20, + "requires": { "features": ["subgroups"] }, + "when": "flashShapeOk", + "constants": { + "headsFromParams": false, + "hasBias": false, + "hasCausal": false, + "combineSubgroups": true, + "usesF16": "qkvUsesF16", + "scalar": "qkvScalar", + "inputVec4": "qkvInputVec4", + "outputScalar": "qkvScalar", + "presentScalar": "qkvCacheScalar", + "inputElement": "qkvInputVec4", + "outputElement": "qkvInputVec4", + "presentElement": "qkvCacheVec4", + "qNumHeads": "attrs.num_heads", + "kvNumHeads": "attrs.kv_num_heads", + "headDim": "dim(shapes.query, 2) / attrs.num_heads", + "headDimV4": "(dim(shapes.query, 2) / attrs.num_heads) / 4", + "qHiddenV4": "dim(shapes.query, 2) / 4", + "kvHidden": "dim(shapes.key, 2)", + "kvHiddenV4": "dim(shapes.key, 2) / 4", + "copyWorkgroupSize": "copyWorkgroupSize", + "presentVec4": true, + "presentElem": "qkvCacheVec4", + "hasWindow": "attrs.local_window_size > 0" + }, + "passes": [ + { + "id": "attention", + "name": "GroupQueryAttention.Flash", + "source": { "shader": "attn-flash-online.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, + "constants": { "usesF16": "tensorDtypes.query == \"float16\"" }, + "bindings": "qkvAttention", + "dispatch": { "x": "dim(shapes.query, 1)", "y": "attrs.num_heads", "z": "dim(shapes.query, 0)" } + }, + { + "id": "present", + "name": "GroupQueryAttention.Present", + "shader": "gqa-present.wgsl.jinja", + "bindings": "qkvPresent", + "dispatch": { "threads": "numel(shapes.presentKeyT) / 4", "workgroupSize": "constants.copyWorkgroupSize" } + } + ] + }, + { + "id": "qkv_present", + "priority": 0, + "when": "fallbackShapeOk", + "constants": { + "headsFromParams": false, + "hasBias": false, + "hasCausal": false, + "hasKeyLimit": false, + "scaleFallbackRsqrt": true, + "usesF16": "qkvUsesF16", + "inputScalar": "qkvScalar", + "outputScalar": "qkvScalar", + "presentScalar": "qkvCacheScalar", + "inputElement": "qkvScalar", + "outputElement": "qkvScalar", + "presentElement": "qkvCacheScalar", + "qNumHeads": "attrs.num_heads", + "kvNumHeads": "attrs.kv_num_heads", + "headDim": "dim(shapes.query, 2) / attrs.num_heads", + "qHidden": "dim(shapes.query, 2)", + "kvHidden": "dim(shapes.key, 2)", + "workgroupSize": "min(256, pow2ceil(dim(shapes.query, 2) / attrs.num_heads))", + "copyWorkgroupSize": "copyWorkgroupSize", + "hasWindow": "attrs.local_window_size > 0" + }, + "passes": [ + { + "id": "attention", + "name": "GroupQueryAttention.Attention", + "source": { "shader": "attn-online-scalar.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, + "constants": { "usesF16": "tensorDtypes.query == \"float16\"" }, + "bindings": "qkvAttention", + "dispatch": { + "x": "min(dim(shapes.query, 1), device.limits.maxComputeWorkgroupsPerDimension)", + "y": "attrs.num_heads", + "z": "dim(shapes.query, 0)" + } + }, + { + "id": "present", + "name": "GroupQueryAttention.Present", + "shader": "gqa-present.wgsl.jinja", + "bindings": "qkvPresent", + "dispatch": { "threads": "numel(shapes.presentKeyT)", "workgroupSize": "constants.copyWorkgroupSize" } + } + ] + }, + { + "id": "quant_int8", + "priority": 35, + "when": ["quantizedPromptOk", "gqaScalarDispatchOk", "attrs.kv_cache_bit_width == 8"], + "constants": { + "headDim": "gqaHeadDim", + "cooperative": "gqaCoop", + "dispatchUnits": "gqaDispatchUnits", + "qHeads": "attrs.num_heads", + "kvHeads": "attrs.kv_num_heads", + "qHidden": "dim(shapes.queryT, 2)", + "packed": "gqaHeadDim / (8 / 8)", + "quantized": true, + "bits": 8, + "mode": "\"build_quant\"", + "qmax": "127 if 8 == 8 else 7", + "qmin": "0 - (128 if 8 == 8 else 8)", + "quantCacheElement": "\"i32\" if 8 == 8 else \"u32\"" + }, + "passes": [ + { + "id": "present", + "name": "GroupQueryAttention.Present", + "shader": "gqa-present.wgsl.jinja", + "bindings": "quantPresent", + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)", + "workgroupSize": "copyWorkgroupSize" + } + }, + { + "id": "attention", + "name": "GroupQueryAttention.Attn", + "shader": "gqa-attention.wgsl.jinja", + "bindings": "quantAttention", + "dispatch": { "workgroups": "constants.dispatchUnits" } + } + ] + }, + { + "id": "quant_int4", + "priority": 35, + "when": ["quantizedPromptOk", "gqaScalarDispatchOk", "attrs.kv_cache_bit_width == 4"], + "constants": { + "headDim": "gqaHeadDim", + "cooperative": "gqaCoop", + "dispatchUnits": "gqaDispatchUnits", + "qHeads": "attrs.num_heads", + "kvHeads": "attrs.kv_num_heads", + "qHidden": "dim(shapes.queryT, 2)", + "packed": "gqaHeadDim / (8 / 4)", + "quantized": true, + "bits": 4, + "mode": "\"build_quant\"", + "qmax": "127 if 4 == 8 else 7", + "qmin": "0 - (128 if 4 == 8 else 8)", + "quantCacheElement": "\"i32\" if 4 == 8 else \"u32\"" + }, + "passes": [ + { + "id": "present", + "name": "GroupQueryAttention.Present", + "shader": "gqa-present.wgsl.jinja", + "bindings": "quantPresent", + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)", + "workgroupSize": "copyWorkgroupSize" + } + }, + { + "id": "attention", + "name": "GroupQueryAttention.Attn", + "shader": "gqa-attention.wgsl.jinja", + "bindings": "quantAttention", + "dispatch": { "workgroups": "constants.dispatchUnits" } + } + ] + }, + { + "id": "quant_int8_decode_splitk", + "description": "Appends int8-quantized key/value rows, partitions cached decode across the key axis, and merges partial online-softmax results. Used when one workgroup per head exposes too little independent work.", + "priority": 45, + "when": ["quantizedCachedDecodeOk", "decodeFlashShapeOk", "cachedNumSplits >= 2", "cachedSplitScratchFits", "subgroupSplitK if \"\" == \"\" else true"], + "constants": { + "headDim": "gqaHeadDim", + "headDimV4": "gqaHeadDim / 4", + "qNumHeads": "attrs.num_heads", + "kvNumHeads": "attrs.kv_num_heads", + "numSplits": "cachedNumSplits", + "usesF16": false, + "scalar": "\"f32\"", + "inputScalar": "\"i32\"", + "inputVec4": "\"vec4\"", + "quantizedCache": true, + "cacheSeqlens": true, + "bits": 8, + "packed": "gqaHeadDim", + "kvHeads": "attrs.kv_num_heads", + "qmax": 127, + "qmin": -128, + "hasBias": false, + "hasWindow": "attrs.local_window_size > 0", + "useSubgroups": "\"\" == \"\"" + }, + "intermediates": [ + { + "id": "partialOut", + "dtype": "float32", + "shape": "[dim(shapes.queryT, 0) * attrs.num_heads * cachedNumSplits * gqaHeadDim]" + }, + { + "id": "partialStats", + "dtype": "float32", + "shape": "[2 * dim(shapes.queryT, 0) * attrs.num_heads * cachedNumSplits]" + } + ], + "passes": [ + { + "id": "copy_past", + "name": "GroupQueryAttention.QuantInt8CopyPast", + "source": { "shader": "gqa-present.wgsl.jinja", "inputs": { "mode": "\"copy\"" } }, + "bindings": "presentCopy", + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", + "workgroupSize": "copyWorkgroupSize" + }, + "viewAlias": [{ "input": "src_k", "output": "present_key" }, { "input": "src_v", "output": "present_value" }] + }, + { + "id": "append", + "name": "GroupQueryAttention.QuantInt8Append", + "source": { "shader": "gqa-present.wgsl.jinja", "inputs": { "mode": "\"append_quant\"" } }, + "bindings": "quantInt8DecodePresent", + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)", + "workgroupSize": "copyWorkgroupSize" + } + }, + { + "id": "split_attention", + "name": "GroupQueryAttention.QuantInt8DecodeSplitK", + "source": { "shader": "attn-flash-decode-splitk.wgsl.jinja", "inputs": { "layout": "\"bhsd\"" } }, + "bindings": "quantInt8SplitAttention", + "dispatch": { "x": "cachedNumSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } + }, + { + "id": "merge", + "name": "GroupQueryAttention.QuantInt8DecodeSplitKMerge", + "source": { "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", "inputs": { "layout": "\"bhsd\"" } }, + "bindings": "splitMergeNoParams", + "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } + } + ] + }, + { + "id": "qkv_present_flash_splitk", + "description": "Partitions direct Q/K/V attention across the key axis, merges partial online-softmax results, and writes the present cache separately. It serves short-query shapes needing more key-axis parallelism.", + "priority": 25, + "when": ["decodeSplitKShapeOk", "qkvSplitScratchFits", "device.features.has(\"subgroups\") if \"\" == \"\" else true"], + "constants": { + "combineSubgroups": "\"\" == \"\"", + "scalar": "qkvScalar", + "inputVec4": "qkvInputVec4", + "outputScalar": "qkvScalar", + "presentScalar": "qkvCacheScalar", + "inputElement": "qkvInputVec4", + "outputElement": "qkvInputVec4", + "presentElement": "qkvCacheVec4", + "qNumHeads": "attrs.num_heads", + "kvNumHeads": "attrs.kv_num_heads", + "headDim": "dim(shapes.query, 2) / attrs.num_heads", + "headDimV4": "(dim(shapes.query, 2) / attrs.num_heads) / 4", + "qHiddenV4": "dim(shapes.query, 2) / 4", + "kvHidden": "dim(shapes.key, 2)", + "kvHiddenV4": "dim(shapes.key, 2) / 4", + "numSplits": "qkvNumSplits", + "copyWorkgroupSize": "copyWorkgroupSize", + "presentVec4": true, + "usesF16": "qkvUsesF16", + "presentElem": "qkvCacheVec4", + "hasWindow": "attrs.local_window_size > 0", + "useSubgroups": "\"\" == \"\"" + }, + "intermediates": [ + { + "id": "partialOut", + "dtype": "float32", + "shape": "[dim(shapes.query, 0) * attrs.num_heads * qkvNumSplits * (dim(shapes.query, 2) / attrs.num_heads)]" + }, + { + "id": "partialStats", + "dtype": "float32", + "shape": "[2 * dim(shapes.query, 0) * attrs.num_heads * qkvNumSplits]" + } + ], + "passes": [ + { + "id": "split_attention", + "name": "GroupQueryAttention.DecodeSplitK", + "source": { "shader": "attn-flash-decode-splitk.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, + "constants": { "usesF16": "tensorDtypes.query == \"float16\"" }, + "bindings": "qkvSplitAttention", + "dispatch": { "x": "qkvNumSplits", "y": "attrs.num_heads", "z": "dim(shapes.query, 0)" } + }, + { + "id": "merge", + "name": "GroupQueryAttention.DecodeSplitKMerge", + "source": { "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, + "constants": { "usesF16": "tensorDtypes.query == \"float16\"" }, + "bindings": "qkvSplitMergeNoParams", + "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.query, 0)" } + }, + { + "id": "present", + "name": "GroupQueryAttention.Present", + "shader": "gqa-present.wgsl.jinja", + "bindings": "qkvPresent", + "dispatch": { "threads": "numel(shapes.presentKeyT) / 4", "workgroupSize": "constants.copyWorkgroupSize" } + } + ] + }, + { + "id": "qkv_present_flash_cluster", + "description": "Computes clustered online-softmax prefill directly from Q/K/V and writes the present cache separately. The family provides subgroup and portable reductions for the same tiled algorithm.", + "priority": 22, + "demoteWhen": ["device.features.has(\"subgroups\") and not qkvF32ClusterRegisterGeometry"], + "when": ["flashShapeOk", "qkvClusterStorageOk if \"\" == \"\" else qkvNoSgClusterStorageOk", "qkvPrefillClusterWorkgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup", "qkvPrefillClusterWorkgroupSize <= device.limits.maxComputeWorkgroupSizeX", "headDim % (4 * qkvPrefillLanesPerQuery) == 0", "qkvF32ClusterRegisterGeometry if \"\" == \"\" else true", "dim(shapes.query, 1) >= tunables.QKV_PREFILL_MIN_QUERY_TOKENS", "subgroupQkvPrefillCluster if \"\" == \"\" else true"], + "constants": { + "hasBias": false, + "hasCausal": false, + "headsFromParams": false, + "combineSubgroups": true, + "usesF16": "qkvUsesF16", + "scalar": "qkvScalar", + "inputVec4": "qkvInputVec4", + "outputScalar": "qkvScalar", + "presentScalar": "qkvCacheScalar", + "inputElement": "qkvInputVec4", + "outputElement": "qkvInputVec4", + "presentElement": "qkvCacheVec4", + "qNumHeads": "attrs.num_heads", + "kvNumHeads": "attrs.kv_num_heads", + "headDim": "dim(shapes.query, 2) / attrs.num_heads", + "headDimV4": "(dim(shapes.query, 2) / attrs.num_heads) / 4", + "qHiddenV4": "dim(shapes.query, 2) / 4", + "kvHidden": "dim(shapes.key, 2)", + "kvHiddenV4": "dim(shapes.key, 2) / 4", + "copyWorkgroupSize": "copyWorkgroupSize", + "presentVec4": true, + "TILE_Q": "qkvPrefillQueryTile", + "TILE_K": "qkvClusterTileK", + "presentElem": "qkvCacheVec4", + "LPQ": "qkvPrefillLanesPerQuery", + "hasWindow": "attrs.local_window_size > 0", + "batchNoSgReduction": "\"\" == \"_nosg\"", + "useSubgroups": "\"\" == \"\"" + }, + "passes": [ + { + "id": "attention", + "name": "GroupQueryAttention.Flash", + "source": { "shader": "attn-flash-prefill-cluster.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, + "constants": { "usesF16": "tensorDtypes.query == \"float16\"" }, + "bindings": "qkvAttention", + "dispatch": { + "x": "ceilDiv(dim(shapes.query, 1), constants.TILE_Q)", + "y": "attrs.num_heads", + "z": "dim(shapes.query, 0)" + } + }, + { + "id": "present", + "name": "GroupQueryAttention.Present", + "shader": "gqa-present.wgsl.jinja", + "bindings": "qkvPresent", + "dispatch": { "threads": "numel(shapes.presentKeyT) / 4", "workgroupSize": "constants.copyWorkgroupSize" } + } + ] + }, + { + "id": "past_kv_bias_flash_prefill", + "priority": 35, + "when": ["sharedKvFloatOk", "prefillFlashShapeOk", "cachedMaskClusterStorageOk if \"\" == \"\" else cachedNoSgClusterStorageOk", "standardSoftmax", "not present.cosCacheT and not rotaryRequested", "not present.qNormWeightT", "present.attentionBiasT", "not present.headSinkT", "(subgroupCachedPrefillCluster and cachedF32ClusterRegisterGeometry) if \"\" == \"\" else true"], + "constants": { + "scalar": "gqaScalar", + "usesF16": "tensorDtypes.queryT == \"float16\"", + "hasCausal": true, + "causalRightAlign": true, + "hasBias": false, + "hasMask": true, + "maskIsBool": false, + "stageMask": "\"\" == \"\"", + "hasWindow": "attrs.local_window_size > 0", + "headDim": "gqaHeadDim", + "headDimV4": "gqaHeadDim / 4", + "qHiddenV4": "dim(shapes.queryT, 2) / 4", + "qNumHeads": "attrs.num_heads", + "kvNumHeads": "attrs.kv_num_heads", + "TILE_Q": "cachedPrefillQueryTile", + "TILE_K": "cachedClusterTileK", + "LPQ": "cachedPrefillLanesPerQuery", + "kvHeads": "attrs.kv_num_heads", + "packed": "gqaHeadDim", + "mode": "\"copy\"", + "inputVec4": "gqaInputVec4", + "inputScalar": "gqaScalar", + "batchNoSgReduction": "\"\" == \"_nosg\"", + "useSubgroups": "\"\" == \"\"" + }, + "passes": [ + { + "id": "present", + "name": "GroupQueryAttention.Present", + "shader": "gqa-present.wgsl.jinja", + "bindings": "presentCopy", + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", + "workgroupSize": "copyWorkgroupSize" + }, + "viewAlias": [{ "input": "src_k", "output": "present_key" }, { "input": "src_v", "output": "present_value" }] + }, + { + "id": "attention", + "name": "GroupQueryAttention.FlashPrefill", + "source": { + "shader": "attn-flash-prefill-cluster.wgsl.jinja", + "inputs": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" } + }, + "bindings": "biasFlashAttention", + "dispatch": { + "x": "ceilDiv(dim(shapes.queryT, 1), constants.TILE_Q)", + "y": "attrs.num_heads", + "z": "dim(shapes.queryT, 0)" + } + } + ] + }, + { + "id": "past_kv_qnorm_rotary_flash_prefill", + "priority": 37, + "when": ["sharedKvFloatOk", "prefillFlashShapeOk", "cachedClusterStorageOk if \"\" == \"\" else cachedNoSgClusterStorageOk", "qPrepScratchFits", "standardSoftmax", "present.cosCacheT and rotaryRequested", "present.qNormWeightT", "not present.attentionBiasT", "not present.headSinkT", "(subgroupCachedPrefillCluster and cachedF32ClusterRegisterGeometry) if \"\" == \"\" else true"], + "constants": { + "scalar": "gqaScalar", + "usesF16": "tensorDtypes.queryT == \"float16\"", + "hasCausal": true, + "causalRightAlign": true, + "hasBias": false, + "hasMask": false, + "maskIsBool": false, + "hasWindow": "attrs.local_window_size > 0", + "headDim": "gqaHeadDim", + "headDimV4": "gqaHeadDim / 4", + "qHiddenV4": "dim(shapes.queryT, 2) / 4", + "qNumHeads": "attrs.num_heads", + "kvNumHeads": "attrs.kv_num_heads", + "TILE_Q": "cachedPrefillQueryTile", + "TILE_K": "cachedClusterTileK", + "LPQ": "cachedPrefillLanesPerQuery", + "kvHeads": "attrs.kv_num_heads", + "packed": "gqaHeadDim", + "mode": "\"copy\"", + "hasRotary": true, + "half": "gqaHeadDim / 2", + "qHeads": "attrs.num_heads", + "qHidden": "dim(shapes.queryT, 2)", + "hasQNorm": true, + "qkEps": "attrs.qk_norm_epsilon", + "inputVec4": "gqaInputVec4", + "inputScalar": "gqaScalar", + "cosScalar": "\"f16\" if tensorDtypes.cosCacheT == \"float16\" else \"f32\"", + "qNormScalar": "\"f16\" if tensorDtypes.qNormWeightT == \"float16\" else \"f32\"", + "batchNoSgReduction": "\"\" == \"_nosg\"", + "useSubgroups": "\"\" == \"\"" + }, + "intermediates": [ + { + "id": "qPrep", + "dtype": "float32", + "shape": "[dim(shapes.queryT, 0), dim(shapes.queryT, 1), dim(shapes.queryT, 2)]" + } + ], + "passes": [ + { + "id": "present", + "name": "GroupQueryAttention.Present", + "shader": "gqa-present.wgsl.jinja", + "bindings": "presentCopy", + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", + "workgroupSize": "copyWorkgroupSize" + }, + "viewAlias": [{ "input": "src_k", "output": "present_key" }, { "input": "src_v", "output": "present_value" }] + }, + { + "id": "qprep", + "name": "GroupQueryAttention.QPrep", + "shader": "gqa-qprep.wgsl.jinja", + "bindings": "qnormRotaryQprep", + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)", + "workgroupSize": "copyWorkgroupSize" + } + }, + { + "id": "attention", + "name": "GroupQueryAttention.FlashPrefill", + "source": { + "shader": "attn-flash-prefill-cluster.wgsl.jinja", + "inputs": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" } + }, + "bindings": "rotaryFlashAttention", + "dispatch": { + "x": "ceilDiv(dim(shapes.queryT, 1), constants.TILE_Q)", + "y": "attrs.num_heads", + "z": "dim(shapes.queryT, 0)" + } + } + ] + }, + { + "id": "quant_int8_decode_splitk_nosg", + "description": "Appends int8-quantized key/value rows, partitions cached decode across the key axis, and merges partial online-softmax results. Used when one workgroup per head exposes too little independent work.", + "priority": 45, + "when": ["quantizedCachedDecodeOk", "decodeFlashShapeOk", "cachedNumSplits >= 2", "cachedSplitScratchFits", "subgroupSplitK if \"_nosg\" == \"\" else true"], + "constants": { + "headDim": "gqaHeadDim", + "headDimV4": "gqaHeadDim / 4", + "qNumHeads": "attrs.num_heads", + "kvNumHeads": "attrs.kv_num_heads", + "numSplits": "cachedNumSplits", + "usesF16": false, + "scalar": "\"f32\"", + "inputScalar": "\"i32\"", + "inputVec4": "\"vec4\"", + "quantizedCache": true, + "cacheSeqlens": true, + "bits": 8, + "packed": "gqaHeadDim", + "kvHeads": "attrs.kv_num_heads", + "qmax": 127, + "qmin": -128, + "hasBias": false, + "hasWindow": "attrs.local_window_size > 0", + "useSubgroups": "\"_nosg\" == \"\"" + }, + "intermediates": [ + { + "id": "partialOut", + "dtype": "float32", + "shape": "[dim(shapes.queryT, 0) * attrs.num_heads * cachedNumSplits * gqaHeadDim]" + }, + { + "id": "partialStats", + "dtype": "float32", + "shape": "[2 * dim(shapes.queryT, 0) * attrs.num_heads * cachedNumSplits]" + } + ], + "passes": [ + { + "id": "copy_past", + "name": "GroupQueryAttention.QuantInt8CopyPast", + "source": { "shader": "gqa-present.wgsl.jinja", "inputs": { "mode": "\"copy\"" } }, + "bindings": "presentCopy", + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", + "workgroupSize": "copyWorkgroupSize" + }, + "viewAlias": [{ "input": "src_k", "output": "present_key" }, { "input": "src_v", "output": "present_value" }] + }, + { + "id": "append", + "name": "GroupQueryAttention.QuantInt8Append", + "source": { "shader": "gqa-present.wgsl.jinja", "inputs": { "mode": "\"append_quant\"" } }, + "bindings": "quantInt8DecodePresent", + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)", + "workgroupSize": "copyWorkgroupSize" + } + }, + { + "id": "split_attention", + "name": "GroupQueryAttention.QuantInt8DecodeSplitK", + "source": { "shader": "attn-flash-decode-splitk.wgsl.jinja", "inputs": { "layout": "\"bhsd\"" } }, + "bindings": "quantInt8SplitAttention", + "dispatch": { "x": "cachedNumSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } + }, + { + "id": "merge", + "name": "GroupQueryAttention.QuantInt8DecodeSplitKMerge", + "source": { "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", "inputs": { "layout": "\"bhsd\"" } }, + "bindings": "splitMergeNoParams", + "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } + } + ] + }, + { + "id": "qkv_present_flash_splitk_nosg", + "description": "Partitions direct Q/K/V attention across the key axis, merges partial online-softmax results, and writes the present cache separately. It serves short-query shapes needing more key-axis parallelism.", + "priority": 25, + "when": ["decodeSplitKShapeOk", "qkvSplitScratchFits", "device.features.has(\"subgroups\") if \"_nosg\" == \"\" else true"], + "constants": { + "combineSubgroups": "\"_nosg\" == \"\"", + "scalar": "qkvScalar", + "inputVec4": "qkvInputVec4", + "outputScalar": "qkvScalar", + "presentScalar": "qkvCacheScalar", + "inputElement": "qkvInputVec4", + "outputElement": "qkvInputVec4", + "presentElement": "qkvCacheVec4", + "qNumHeads": "attrs.num_heads", + "kvNumHeads": "attrs.kv_num_heads", + "headDim": "dim(shapes.query, 2) / attrs.num_heads", + "headDimV4": "(dim(shapes.query, 2) / attrs.num_heads) / 4", + "qHiddenV4": "dim(shapes.query, 2) / 4", + "kvHidden": "dim(shapes.key, 2)", + "kvHiddenV4": "dim(shapes.key, 2) / 4", + "numSplits": "qkvNumSplits", + "copyWorkgroupSize": "copyWorkgroupSize", + "presentVec4": true, + "usesF16": "qkvUsesF16", + "presentElem": "qkvCacheVec4", + "hasWindow": "attrs.local_window_size > 0", + "useSubgroups": "\"_nosg\" == \"\"" + }, + "intermediates": [ + { + "id": "partialOut", + "dtype": "float32", + "shape": "[dim(shapes.query, 0) * attrs.num_heads * qkvNumSplits * (dim(shapes.query, 2) / attrs.num_heads)]" + }, + { + "id": "partialStats", + "dtype": "float32", + "shape": "[2 * dim(shapes.query, 0) * attrs.num_heads * qkvNumSplits]" + } + ], + "passes": [ + { + "id": "split_attention", + "name": "GroupQueryAttention.DecodeSplitK", + "source": { "shader": "attn-flash-decode-splitk.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, + "constants": { "usesF16": "tensorDtypes.query == \"float16\"" }, + "bindings": "qkvSplitAttention", + "dispatch": { "x": "qkvNumSplits", "y": "attrs.num_heads", "z": "dim(shapes.query, 0)" } + }, + { + "id": "merge", + "name": "GroupQueryAttention.DecodeSplitKMerge", + "source": { "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, + "constants": { "usesF16": "tensorDtypes.query == \"float16\"" }, + "bindings": "qkvSplitMergeNoParams", + "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.query, 0)" } + }, + { + "id": "present", + "name": "GroupQueryAttention.Present", + "shader": "gqa-present.wgsl.jinja", + "bindings": "qkvPresent", + "dispatch": { "threads": "numel(shapes.presentKeyT) / 4", "workgroupSize": "constants.copyWorkgroupSize" } + } + ] + }, + { + "id": "qkv_present_flash_cluster_nosg", + "description": "Computes clustered online-softmax prefill directly from Q/K/V and writes the present cache separately. The family provides subgroup and portable reductions for the same tiled algorithm.", + "priority": 22, + "demoteWhen": ["device.features.has(\"subgroups\") and not qkvF32ClusterRegisterGeometry"], + "when": ["flashShapeOk", "qkvClusterStorageOk if \"_nosg\" == \"\" else qkvNoSgClusterStorageOk", "qkvPrefillClusterWorkgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup", "qkvPrefillClusterWorkgroupSize <= device.limits.maxComputeWorkgroupSizeX", "headDim % (4 * qkvPrefillLanesPerQuery) == 0", "qkvF32ClusterRegisterGeometry if \"_nosg\" == \"\" else true", "dim(shapes.query, 1) >= tunables.QKV_PREFILL_MIN_QUERY_TOKENS", "subgroupQkvPrefillCluster if \"_nosg\" == \"\" else true"], + "constants": { + "hasBias": false, + "hasCausal": false, + "headsFromParams": false, + "combineSubgroups": true, + "usesF16": "qkvUsesF16", + "scalar": "qkvScalar", + "inputVec4": "qkvInputVec4", + "outputScalar": "qkvScalar", + "presentScalar": "qkvCacheScalar", + "inputElement": "qkvInputVec4", + "outputElement": "qkvInputVec4", + "presentElement": "qkvCacheVec4", + "qNumHeads": "attrs.num_heads", + "kvNumHeads": "attrs.kv_num_heads", + "headDim": "dim(shapes.query, 2) / attrs.num_heads", + "headDimV4": "(dim(shapes.query, 2) / attrs.num_heads) / 4", + "qHiddenV4": "dim(shapes.query, 2) / 4", + "kvHidden": "dim(shapes.key, 2)", + "kvHiddenV4": "dim(shapes.key, 2) / 4", + "copyWorkgroupSize": "copyWorkgroupSize", + "presentVec4": true, + "TILE_Q": "qkvPrefillQueryTile", + "TILE_K": "qkvNoSgTileK", + "presentElem": "qkvCacheVec4", + "LPQ": "qkvPrefillLanesPerQuery", + "hasWindow": "attrs.local_window_size > 0", + "batchNoSgReduction": "\"_nosg\" == \"_nosg\"", + "useSubgroups": "\"_nosg\" == \"\"" + }, + "passes": [ + { + "id": "attention", + "name": "GroupQueryAttention.Flash", + "source": { "shader": "attn-flash-prefill-cluster.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, + "constants": { "usesF16": "tensorDtypes.query == \"float16\"" }, + "bindings": "qkvAttention", + "dispatch": { + "x": "ceilDiv(dim(shapes.query, 1), constants.TILE_Q)", + "y": "attrs.num_heads", + "z": "dim(shapes.query, 0)" + } + }, + { + "id": "present", + "name": "GroupQueryAttention.Present", + "shader": "gqa-present.wgsl.jinja", + "bindings": "qkvPresent", + "dispatch": { "threads": "numel(shapes.presentKeyT) / 4", "workgroupSize": "constants.copyWorkgroupSize" } + } + ] + }, + { + "id": "past_kv_bias_flash_prefill_nosg", + "priority": 35, + "when": ["sharedKvFloatOk", "prefillFlashShapeOk", "cachedMaskClusterStorageOk if \"_nosg\" == \"\" else cachedNoSgClusterStorageOk", "standardSoftmax", "not present.cosCacheT and not rotaryRequested", "not present.qNormWeightT", "present.attentionBiasT", "not present.headSinkT", "(subgroupCachedPrefillCluster and cachedF32ClusterRegisterGeometry) if \"_nosg\" == \"\" else true"], + "constants": { + "scalar": "gqaScalar", + "usesF16": "tensorDtypes.queryT == \"float16\"", + "hasCausal": true, + "causalRightAlign": true, + "hasBias": false, + "hasMask": true, + "maskIsBool": false, + "stageMask": "\"_nosg\" == \"\"", + "hasWindow": "attrs.local_window_size > 0", + "headDim": "gqaHeadDim", + "headDimV4": "gqaHeadDim / 4", + "qHiddenV4": "dim(shapes.queryT, 2) / 4", + "qNumHeads": "attrs.num_heads", + "kvNumHeads": "attrs.kv_num_heads", + "TILE_Q": "cachedPrefillQueryTile", + "TILE_K": "cachedNoSgTileK", + "LPQ": "cachedPrefillLanesPerQuery", + "kvHeads": "attrs.kv_num_heads", + "packed": "gqaHeadDim", + "mode": "\"copy\"", + "inputVec4": "gqaInputVec4", + "inputScalar": "gqaScalar", + "batchNoSgReduction": "\"_nosg\" == \"_nosg\"", + "useSubgroups": "\"_nosg\" == \"\"" + }, + "passes": [ + { + "id": "present", + "name": "GroupQueryAttention.Present", + "shader": "gqa-present.wgsl.jinja", + "bindings": "presentCopy", + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", + "workgroupSize": "copyWorkgroupSize" + }, + "viewAlias": [{ "input": "src_k", "output": "present_key" }, { "input": "src_v", "output": "present_value" }] + }, + { + "id": "attention", + "name": "GroupQueryAttention.FlashPrefill", + "source": { + "shader": "attn-flash-prefill-cluster.wgsl.jinja", + "inputs": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" } + }, + "bindings": "biasFlashAttention", + "dispatch": { + "x": "ceilDiv(dim(shapes.queryT, 1), constants.TILE_Q)", + "y": "attrs.num_heads", + "z": "dim(shapes.queryT, 0)" + } + } + ] + }, + { + "id": "past_kv_qnorm_rotary_flash_prefill_nosg", + "priority": 37, + "when": ["sharedKvFloatOk", "prefillFlashShapeOk", "cachedClusterStorageOk if \"_nosg\" == \"\" else cachedNoSgClusterStorageOk", "qPrepScratchFits", "standardSoftmax", "present.cosCacheT and rotaryRequested", "present.qNormWeightT", "not present.attentionBiasT", "not present.headSinkT", "(subgroupCachedPrefillCluster and cachedF32ClusterRegisterGeometry) if \"_nosg\" == \"\" else true"], + "constants": { + "scalar": "gqaScalar", + "usesF16": "tensorDtypes.queryT == \"float16\"", + "hasCausal": true, + "causalRightAlign": true, + "hasBias": false, + "hasMask": false, + "maskIsBool": false, + "hasWindow": "attrs.local_window_size > 0", + "headDim": "gqaHeadDim", + "headDimV4": "gqaHeadDim / 4", + "qHiddenV4": "dim(shapes.queryT, 2) / 4", + "qNumHeads": "attrs.num_heads", + "kvNumHeads": "attrs.kv_num_heads", + "TILE_Q": "cachedPrefillQueryTile", + "TILE_K": "cachedNoSgTileK", + "LPQ": "cachedPrefillLanesPerQuery", + "kvHeads": "attrs.kv_num_heads", + "packed": "gqaHeadDim", + "mode": "\"copy\"", + "hasRotary": true, + "half": "gqaHeadDim / 2", + "qHeads": "attrs.num_heads", + "qHidden": "dim(shapes.queryT, 2)", + "hasQNorm": true, + "qkEps": "attrs.qk_norm_epsilon", + "inputVec4": "gqaInputVec4", + "inputScalar": "gqaScalar", + "cosScalar": "\"f16\" if tensorDtypes.cosCacheT == \"float16\" else \"f32\"", + "qNormScalar": "\"f16\" if tensorDtypes.qNormWeightT == \"float16\" else \"f32\"", + "batchNoSgReduction": "\"_nosg\" == \"_nosg\"", + "useSubgroups": "\"_nosg\" == \"\"" + }, + "intermediates": [ + { + "id": "qPrep", + "dtype": "float32", + "shape": "[dim(shapes.queryT, 0), dim(shapes.queryT, 1), dim(shapes.queryT, 2)]" + } + ], + "passes": [ + { + "id": "present", + "name": "GroupQueryAttention.Present", + "shader": "gqa-present.wgsl.jinja", + "bindings": "presentCopy", + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", + "workgroupSize": "copyWorkgroupSize" + }, + "viewAlias": [{ "input": "src_k", "output": "present_key" }, { "input": "src_v", "output": "present_value" }] + }, + { + "id": "qprep", + "name": "GroupQueryAttention.QPrep", + "shader": "gqa-qprep.wgsl.jinja", + "bindings": "qnormRotaryQprep", + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)", + "workgroupSize": "copyWorkgroupSize" + } + }, + { + "id": "attention", + "name": "GroupQueryAttention.FlashPrefill", + "source": { + "shader": "attn-flash-prefill-cluster.wgsl.jinja", + "inputs": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" } + }, + "bindings": "rotaryFlashAttention", + "dispatch": { + "x": "ceilDiv(dim(shapes.queryT, 1), constants.TILE_Q)", + "y": "attrs.num_heads", + "z": "dim(shapes.queryT, 0)" + } + } + ] + }, + { + "id": "past_kv_decode_splitk", + "description": "Updates the float cache in copy, merge, or window-shift mode, then partitions cached decode across the key axis and combines partial online-softmax results. The family includes subgroup and portable forms.", + "priority": 38, + "when": ["sharedKvFloatOk", "decodeFlashShapeOk", "cachedSplitScratchFits", "plainAttentionOptions", "subgroupSplitK if \"\" == \"\" else true"], + "constants": { + "cacheSeqlens": true, + "scalar": "gqaScalar", + "usesF16": "tensorDtypes.queryT == \"float16\"", + "qNumHeads": "attrs.num_heads", + "kvNumHeads": "attrs.kv_num_heads", + "headDim": "gqaHeadDim", + "headDimV4": "gqaHeadDim / 4", + "numSplits": "cachedNumSplits", + "hasWindow": "attrs.local_window_size > 0", + "kvHeads": "attrs.kv_num_heads", + "packed": "gqaHeadDim", + "mode": "\"copy\"", + "inputVec4": "gqaInputVec4", + "inputScalar": "gqaScalar", + "zeroScalar": "gqaScalar", + "combineSubgroups": "\"\" == \"\"", + "useSubgroups": "\"\" == \"\"" + }, + "intermediates": [ + { + "id": "partialOut", + "dtype": "float32", + "shape": "[dim(shapes.queryT, 0) * attrs.num_heads * (cachedNumSplits) * gqaHeadDim]" + }, + { + "id": "partialStats", + "dtype": "float32", + "shape": "[2 * dim(shapes.queryT, 0) * attrs.num_heads * (cachedNumSplits)]" + } + ], + "passes": [ + { + "id": "present", + "name": "GroupQueryAttention.Present", + "shader": "gqa-present.wgsl.jinja", + "bindings": "presentCopy", + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", + "workgroupSize": "copyWorkgroupSize" + }, + "viewAlias": [{ "input": "src_k", "output": "present_key" }, { "input": "src_v", "output": "present_value" }] + }, + { + "id": "split_attention", + "name": "GroupQueryAttention.DecodeSplitK", + "source": { "shader": "attn-flash-decode-splitk.wgsl.jinja", "inputs": { "layout": "\"bhsd\"" } }, + "bindings": "splitSeqlensAttention", + "dispatch": { "x": "cachedNumSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } + }, + { + "id": "merge", + "name": "GroupQueryAttention.DecodeSplitKMerge", + "source": { "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", "inputs": { "layout": "\"bhsd\"" } }, + "bindings": "splitMergeNoParams", + "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } + } + ] + }, + { + "id": "new_kv_past_decode_splitk", + "description": "Updates the float cache in copy, merge, or window-shift mode, then partitions cached decode across the key axis and combines partial online-softmax results. The family includes subgroup and portable forms.", + "priority": 38, + "when": ["newKvPastOk", "decodeFlashShapeOk", "cachedSplitScratchFits", "plainAttentionOptions", "subgroupSplitK if \"\" == \"\" else true"], + "constants": { + "scalar": "gqaScalar", + "usesF16": "tensorDtypes.queryT == \"float16\"", + "qNumHeads": "attrs.num_heads", + "kvNumHeads": "attrs.kv_num_heads", + "headDim": "gqaHeadDim", + "headDimV4": "gqaHeadDim / 4", + "numSplits": "cachedNumSplits", + "hasWindow": "attrs.local_window_size > 0", + "kvHeads": "attrs.kv_num_heads", + "packed": "gqaHeadDim", + "mode": "\"merge\"", + "inputVec4": "gqaInputVec4", + "inputScalar": "gqaScalar", + "zeroScalar": "gqaScalar", + "combineSubgroups": "\"\" == \"\"", + "useSubgroups": "\"\" == \"\"" + }, + "intermediates": [ + { + "id": "partialOut", + "dtype": "float32", + "shape": "[dim(shapes.queryT, 0) * attrs.num_heads * (cachedNumSplits) * gqaHeadDim]" + }, + { + "id": "partialStats", + "dtype": "float32", + "shape": "[2 * dim(shapes.queryT, 0) * attrs.num_heads * (cachedNumSplits)]" + } + ], + "passes": [ + { + "id": "present", + "name": "GroupQueryAttention.Merge", + "shader": "gqa-present.wgsl.jinja", + "bindings": "presentMerge", + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)", + "workgroupSize": "copyWorkgroupSize" + } + }, + { + "id": "split_attention", + "name": "GroupQueryAttention.DecodeSplitK", + "source": { "shader": "attn-flash-decode-splitk.wgsl.jinja", "inputs": { "layout": "\"bhsd\"" } }, + "bindings": "splitAttention", + "dispatch": { "x": "cachedNumSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } + }, + { + "id": "merge", + "name": "GroupQueryAttention.DecodeSplitKMerge", + "source": { "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", "inputs": { "layout": "\"bhsd\"" } }, + "bindings": "splitMergeNoParams", + "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } + } + ] + }, + { + "id": "window_shift_decode_splitk", + "description": "Updates the float cache in copy, merge, or window-shift mode, then partitions cached decode across the key axis and combines partial online-softmax results. The family includes subgroup and portable forms.", + "priority": 38, + "when": ["windowShiftOk", "decodeFlashShapeOk", "cachedSplitScratchFits", "plainAttentionOptions", "subgroupSplitK if \"\" == \"\" else true"], + "constants": { + "cacheSeqlens": true, + "scalar": "gqaScalar", + "usesF16": "tensorDtypes.queryT == \"float16\"", + "qNumHeads": "attrs.num_heads", + "kvNumHeads": "attrs.kv_num_heads", + "headDim": "gqaHeadDim", + "headDimV4": "gqaHeadDim / 4", + "numSplits": "cachedNumSplits", + "hasWindow": "attrs.local_window_size > 0", + "kvHeads": "attrs.kv_num_heads", + "packed": "gqaHeadDim", + "mode": "\"window_shift\"", + "inputVec4": "gqaInputVec4", + "inputScalar": "gqaScalar", + "zeroScalar": "gqaScalar", + "combineSubgroups": "\"\" == \"\"", + "useSubgroups": "\"\" == \"\"" + }, + "intermediates": [ + { + "id": "partialOut", + "dtype": "float32", + "shape": "[dim(shapes.queryT, 0) * attrs.num_heads * (cachedNumSplits) * gqaHeadDim]" + }, + { + "id": "partialStats", + "dtype": "float32", + "shape": "[2 * dim(shapes.queryT, 0) * attrs.num_heads * (cachedNumSplits)]" + } + ], + "passes": [ + { + "id": "present", + "name": "GroupQueryAttention.WindowShift", + "shader": "gqa-present.wgsl.jinja", + "bindings": "presentMergeShare", + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)", + "workgroupSize": "copyWorkgroupSize" + } + }, + { + "id": "split_attention", + "name": "GroupQueryAttention.DecodeSplitK", + "source": { "shader": "attn-flash-decode-splitk.wgsl.jinja", "inputs": { "layout": "\"bhsd\"" } }, + "bindings": "splitSeqlensAttention", + "dispatch": { "x": "cachedNumSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } + }, + { + "id": "merge", + "name": "GroupQueryAttention.DecodeSplitKMerge", + "source": { "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", "inputs": { "layout": "\"bhsd\"" } }, + "bindings": "splitMergeNoParams", + "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } + } + ] + }, + { + "id": "past_kv_decode_splitk_nosg", + "description": "Updates the float cache in copy, merge, or window-shift mode, then partitions cached decode across the key axis and combines partial online-softmax results. The family includes subgroup and portable forms.", + "priority": 38, + "when": ["sharedKvFloatOk", "decodeFlashShapeOk", "cachedSplitScratchFits", "plainAttentionOptions", "subgroupSplitK if \"_nosg\" == \"\" else true"], + "constants": { + "cacheSeqlens": true, + "scalar": "gqaScalar", + "usesF16": "tensorDtypes.queryT == \"float16\"", + "qNumHeads": "attrs.num_heads", + "kvNumHeads": "attrs.kv_num_heads", + "headDim": "gqaHeadDim", + "headDimV4": "gqaHeadDim / 4", + "numSplits": "cachedNumSplits", + "hasWindow": "attrs.local_window_size > 0", + "kvHeads": "attrs.kv_num_heads", + "packed": "gqaHeadDim", + "mode": "\"copy\"", + "inputVec4": "gqaInputVec4", + "inputScalar": "gqaScalar", + "zeroScalar": "gqaScalar", + "combineSubgroups": "\"_nosg\" == \"\"", + "useSubgroups": "\"_nosg\" == \"\"" + }, + "intermediates": [ + { + "id": "partialOut", + "dtype": "float32", + "shape": "[dim(shapes.queryT, 0) * attrs.num_heads * (cachedNumSplits) * gqaHeadDim]" + }, + { + "id": "partialStats", + "dtype": "float32", + "shape": "[2 * dim(shapes.queryT, 0) * attrs.num_heads * (cachedNumSplits)]" + } + ], + "passes": [ + { + "id": "present", + "name": "GroupQueryAttention.Present", + "shader": "gqa-present.wgsl.jinja", + "bindings": "presentCopy", + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", + "workgroupSize": "copyWorkgroupSize" + }, + "viewAlias": [{ "input": "src_k", "output": "present_key" }, { "input": "src_v", "output": "present_value" }] + }, + { + "id": "split_attention", + "name": "GroupQueryAttention.DecodeSplitK", + "source": { "shader": "attn-flash-decode-splitk.wgsl.jinja", "inputs": { "layout": "\"bhsd\"" } }, + "bindings": "splitSeqlensAttention", + "dispatch": { "x": "cachedNumSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } + }, + { + "id": "merge", + "name": "GroupQueryAttention.DecodeSplitKMerge", + "source": { "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", "inputs": { "layout": "\"bhsd\"" } }, + "bindings": "splitMergeNoParams", + "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } + } + ] + }, + { + "id": "new_kv_past_decode_splitk_nosg", + "description": "Updates the float cache in copy, merge, or window-shift mode, then partitions cached decode across the key axis and combines partial online-softmax results. The family includes subgroup and portable forms.", + "priority": 38, + "when": ["newKvPastOk", "decodeFlashShapeOk", "cachedSplitScratchFits", "plainAttentionOptions", "subgroupSplitK if \"_nosg\" == \"\" else true"], + "constants": { + "scalar": "gqaScalar", + "usesF16": "tensorDtypes.queryT == \"float16\"", + "qNumHeads": "attrs.num_heads", + "kvNumHeads": "attrs.kv_num_heads", + "headDim": "gqaHeadDim", + "headDimV4": "gqaHeadDim / 4", + "numSplits": "cachedNumSplits", + "hasWindow": "attrs.local_window_size > 0", + "kvHeads": "attrs.kv_num_heads", + "packed": "gqaHeadDim", + "mode": "\"merge\"", + "inputVec4": "gqaInputVec4", + "inputScalar": "gqaScalar", + "zeroScalar": "gqaScalar", + "combineSubgroups": "\"_nosg\" == \"\"", + "useSubgroups": "\"_nosg\" == \"\"" + }, + "intermediates": [ + { + "id": "partialOut", + "dtype": "float32", + "shape": "[dim(shapes.queryT, 0) * attrs.num_heads * (cachedNumSplits) * gqaHeadDim]" + }, + { + "id": "partialStats", + "dtype": "float32", + "shape": "[2 * dim(shapes.queryT, 0) * attrs.num_heads * (cachedNumSplits)]" + } + ], + "passes": [ + { + "id": "present", + "name": "GroupQueryAttention.Merge", + "shader": "gqa-present.wgsl.jinja", + "bindings": "presentMerge", + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)", + "workgroupSize": "copyWorkgroupSize" + } + }, + { + "id": "split_attention", + "name": "GroupQueryAttention.DecodeSplitK", + "source": { "shader": "attn-flash-decode-splitk.wgsl.jinja", "inputs": { "layout": "\"bhsd\"" } }, + "bindings": "splitAttention", + "dispatch": { "x": "cachedNumSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } + }, + { + "id": "merge", + "name": "GroupQueryAttention.DecodeSplitKMerge", + "source": { "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", "inputs": { "layout": "\"bhsd\"" } }, + "bindings": "splitMergeNoParams", + "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } + } + ] + }, + { + "id": "window_shift_decode_splitk_nosg", + "description": "Updates the float cache in copy, merge, or window-shift mode, then partitions cached decode across the key axis and combines partial online-softmax results. The family includes subgroup and portable forms.", + "priority": 38, + "when": ["windowShiftOk", "decodeFlashShapeOk", "cachedSplitScratchFits", "plainAttentionOptions", "subgroupSplitK if \"_nosg\" == \"\" else true"], + "constants": { + "cacheSeqlens": true, + "scalar": "gqaScalar", + "usesF16": "tensorDtypes.queryT == \"float16\"", + "qNumHeads": "attrs.num_heads", + "kvNumHeads": "attrs.kv_num_heads", + "headDim": "gqaHeadDim", + "headDimV4": "gqaHeadDim / 4", + "numSplits": "cachedNumSplits", + "hasWindow": "attrs.local_window_size > 0", + "kvHeads": "attrs.kv_num_heads", + "packed": "gqaHeadDim", + "mode": "\"window_shift\"", + "inputVec4": "gqaInputVec4", + "inputScalar": "gqaScalar", + "zeroScalar": "gqaScalar", + "combineSubgroups": "\"_nosg\" == \"\"", + "useSubgroups": "\"_nosg\" == \"\"" + }, + "intermediates": [ + { + "id": "partialOut", + "dtype": "float32", + "shape": "[dim(shapes.queryT, 0) * attrs.num_heads * (cachedNumSplits) * gqaHeadDim]" + }, + { + "id": "partialStats", + "dtype": "float32", + "shape": "[2 * dim(shapes.queryT, 0) * attrs.num_heads * (cachedNumSplits)]" + } + ], + "passes": [ + { + "id": "present", + "name": "GroupQueryAttention.WindowShift", + "shader": "gqa-present.wgsl.jinja", + "bindings": "presentMergeShare", + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)", + "workgroupSize": "copyWorkgroupSize" + } + }, + { + "id": "split_attention", + "name": "GroupQueryAttention.DecodeSplitK", + "source": { "shader": "attn-flash-decode-splitk.wgsl.jinja", "inputs": { "layout": "\"bhsd\"" } }, + "bindings": "splitSeqlensAttention", + "dispatch": { "x": "cachedNumSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } + }, + { + "id": "merge", + "name": "GroupQueryAttention.DecodeSplitKMerge", + "source": { "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", "inputs": { "layout": "\"bhsd\"" } }, + "bindings": "splitMergeNoParams", + "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } + } + ] + }, + { + "id": "past_kv_flash_prefill", + "description": "Updates the float cache in the requested mode and applies clustered online-softmax causal prefill. Cases cover cache merging and optional attention semantics in subgroup and portable forms.", + "priority": 34, + "when": ["sharedKvFloatOk", "prefillFlashShapeOk", "cachedClusterStorageOk if \"\" == \"\" else cachedNoSgClusterStorageOk", "plainAttentionOptions", "(subgroupCachedPrefillCluster and cachedF32ClusterRegisterGeometry) if \"\" == \"\" else true"], + "constants": { + "scalar": "gqaScalar", + "usesF16": "tensorDtypes.queryT == \"float16\"", + "hasCausal": true, + "causalRightAlign": true, + "hasBias": false, + "hasMask": false, + "maskIsBool": false, + "hasWindow": "attrs.local_window_size > 0", + "headDim": "gqaHeadDim", + "headDimV4": "gqaHeadDim / 4", + "qHiddenV4": "dim(shapes.queryT, 2) / 4", + "qNumHeads": "attrs.num_heads", + "kvNumHeads": "attrs.kv_num_heads", + "TILE_Q": "cachedPrefillQueryTile", + "TILE_K": "cachedClusterTileK", + "LPQ": "cachedPrefillLanesPerQuery", + "kvHeads": "attrs.kv_num_heads", + "packed": "gqaHeadDim", + "mode": "\"copy\"", + "inputVec4": "gqaInputVec4", + "inputScalar": "gqaScalar", + "batchNoSgReduction": "\"\" == \"_nosg\"", + "useSubgroups": "\"\" == \"\"" + }, + "passes": [ + { + "id": "present", + "name": "GroupQueryAttention.Present", + "shader": "gqa-present.wgsl.jinja", + "bindings": "presentCopy", + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", + "workgroupSize": "copyWorkgroupSize" + }, + "viewAlias": [{ "input": "src_k", "output": "present_key" }, { "input": "src_v", "output": "present_value" }] + }, + { + "id": "attention", + "name": "GroupQueryAttention.FlashPrefill", + "source": { + "shader": "attn-flash-prefill-cluster.wgsl.jinja", + "inputs": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" } + }, + "bindings": "flashPrefillPast", + "dispatch": { + "x": "ceilDiv(dim(shapes.queryT, 1), constants.TILE_Q)", + "y": "attrs.num_heads", + "z": "dim(shapes.queryT, 0)" + } + } + ] + }, + { + "id": "new_kv_past_flash_prefill", + "description": "Updates the float cache in the requested mode and applies clustered online-softmax causal prefill. Cases cover cache merging and optional attention semantics in subgroup and portable forms.", + "priority": 34, + "when": ["newKvPastOk", "prefillFlashShapeOk", "cachedClusterStorageOk if \"\" == \"\" else cachedNoSgClusterStorageOk", "plainAttentionOptions", "(subgroupCachedPrefillCluster and cachedF32ClusterRegisterGeometry) if \"\" == \"\" else true"], + "constants": { + "scalar": "gqaScalar", + "usesF16": "tensorDtypes.queryT == \"float16\"", + "hasCausal": true, + "causalRightAlign": true, + "hasBias": false, + "hasMask": false, + "maskIsBool": false, + "hasWindow": "attrs.local_window_size > 0", + "headDim": "gqaHeadDim", + "headDimV4": "gqaHeadDim / 4", + "qHiddenV4": "dim(shapes.queryT, 2) / 4", + "qNumHeads": "attrs.num_heads", + "kvNumHeads": "attrs.kv_num_heads", + "TILE_Q": "cachedPrefillQueryTile", + "TILE_K": "cachedClusterTileK", + "LPQ": "cachedPrefillLanesPerQuery", + "kvHeads": "attrs.kv_num_heads", + "packed": "gqaHeadDim", + "mode": "\"merge\"", + "inputVec4": "gqaInputVec4", + "inputScalar": "gqaScalar", + "batchNoSgReduction": "\"\" == \"_nosg\"", + "useSubgroups": "\"\" == \"\"" + }, + "passes": [ + { + "id": "present", + "name": "GroupQueryAttention.Merge", + "shader": "gqa-present.wgsl.jinja", + "bindings": "presentMerge", + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)", + "workgroupSize": "copyWorkgroupSize" + } + }, + { + "id": "attention", + "name": "GroupQueryAttention.FlashPrefill", + "source": { + "shader": "attn-flash-prefill-cluster.wgsl.jinja", + "inputs": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" } + }, + "bindings": "flashPrefillNew", + "dispatch": { + "x": "ceilDiv(dim(shapes.queryT, 1), constants.TILE_Q)", + "y": "attrs.num_heads", + "z": "dim(shapes.queryT, 0)" + } + } + ] + }, + { + "id": "past_kv_rotary_flash_prefill", + "description": "Updates the float cache in the requested mode and applies clustered online-softmax causal prefill. Cases cover cache merging and optional attention semantics in subgroup and portable forms.", + "priority": 36, + "when": ["sharedKvFloatOk", "prefillFlashShapeOk", "cachedClusterStorageOk if \"\" == \"\" else cachedNoSgClusterStorageOk", "standardSoftmax", "present.cosCacheT and rotaryRequested", "not present.qNormWeightT", "not present.attentionBiasT", "not present.headSinkT", "(subgroupCachedPrefillCluster and cachedF32ClusterRegisterGeometry) if \"\" == \"\" else true"], + "constants": { + "scalar": "gqaScalar", + "usesF16": "tensorDtypes.queryT == \"float16\"", + "hasCausal": true, + "causalRightAlign": true, + "hasBias": false, + "hasMask": false, + "maskIsBool": false, + "hasWindow": "attrs.local_window_size > 0", + "headDim": "gqaHeadDim", + "headDimV4": "gqaHeadDim / 4", + "qHiddenV4": "dim(shapes.queryT, 2) / 4", + "qNumHeads": "attrs.num_heads", + "kvNumHeads": "attrs.kv_num_heads", + "TILE_Q": "cachedPrefillQueryTile", + "TILE_K": "cachedClusterTileK", + "LPQ": "cachedPrefillLanesPerQuery", + "kvHeads": "attrs.kv_num_heads", + "packed": "gqaHeadDim", + "mode": "\"copy\"", + "hasRotary": true, + "fusedRotary": true, + "half": "gqaHeadDim / 2", + "qHeads": "attrs.num_heads", + "qHidden": "dim(shapes.queryT, 2)", + "inputVec4": "gqaInputVec4", + "inputScalar": "gqaScalar", + "cosScalar": "\"f16\" if tensorDtypes.cosCacheT == \"float16\" else \"f32\"", + "batchNoSgReduction": "\"\" == \"_nosg\"", + "useSubgroups": "\"\" == \"\"" + }, + "passes": [ + { + "id": "present", + "name": "GroupQueryAttention.Present", + "shader": "gqa-present.wgsl.jinja", + "bindings": "presentCopy", + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", + "workgroupSize": "copyWorkgroupSize" + }, + "viewAlias": [{ "input": "src_k", "output": "present_key" }, { "input": "src_v", "output": "present_value" }] + }, + { + "id": "attention", + "name": "GroupQueryAttention.FlashPrefill", + "source": { + "shader": "attn-flash-prefill-cluster.wgsl.jinja", + "inputs": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" } + }, + "bindings": "rotaryFusedFlashAttention", + "dispatch": { + "x": "ceilDiv(dim(shapes.queryT, 1), constants.TILE_Q)", + "y": "attrs.num_heads", + "z": "dim(shapes.queryT, 0)" + } + } + ] + }, + { + "id": "past_kv_softcap_flash_prefill", + "description": "Updates the float cache in the requested mode and applies clustered online-softmax causal prefill. Cases cover cache merging and optional attention semantics in subgroup and portable forms.", + "priority": 34, + "when": ["sharedKvFloatOk", "prefillFlashShapeOk", "cachedClusterStorageOk if \"\" == \"\" else cachedNoSgClusterStorageOk", "noAuxAttentionInputs", "attrs.softcap > 0", "attrs.smooth_softmax != 1", "(subgroupCachedPrefillCluster and cachedF32ClusterRegisterGeometry) if \"\" == \"\" else true"], + "constants": { + "scalar": "gqaScalar", + "usesF16": "tensorDtypes.queryT == \"float16\"", + "hasCausal": true, + "causalRightAlign": true, + "hasBias": false, + "hasSoftcap": true, + "hasMask": false, + "maskIsBool": false, + "hasWindow": "attrs.local_window_size > 0", + "headDim": "gqaHeadDim", + "headDimV4": "gqaHeadDim / 4", + "qHiddenV4": "dim(shapes.queryT, 2) / 4", + "qNumHeads": "attrs.num_heads", + "kvNumHeads": "attrs.kv_num_heads", + "TILE_Q": "cachedPrefillQueryTile", + "TILE_K": "cachedClusterTileK", + "LPQ": "cachedPrefillLanesPerQuery", + "kvHeads": "attrs.kv_num_heads", + "packed": "gqaHeadDim", + "mode": "\"copy\"", + "inputVec4": "gqaInputVec4", + "inputScalar": "gqaScalar", + "batchNoSgReduction": "\"\" == \"_nosg\"", + "useSubgroups": "\"\" == \"\"" + }, + "passes": [ + { + "id": "present", + "name": "GroupQueryAttention.Present", + "shader": "gqa-present.wgsl.jinja", + "bindings": "presentCopy", + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", + "workgroupSize": "copyWorkgroupSize" + }, + "viewAlias": [{ "input": "src_k", "output": "present_key" }, { "input": "src_v", "output": "present_value" }] + }, + { + "id": "attention", + "name": "GroupQueryAttention.FlashPrefill", + "source": { + "shader": "attn-flash-prefill-cluster.wgsl.jinja", + "inputs": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" } + }, + "bindings": "flashPrefillPastSoftcap", + "dispatch": { + "x": "ceilDiv(dim(shapes.queryT, 1), constants.TILE_Q)", + "y": "attrs.num_heads", + "z": "dim(shapes.queryT, 0)" + } + } + ] + }, + { + "id": "past_kv_headsink_flash_prefill", + "description": "Updates the float cache in the requested mode and applies clustered online-softmax causal prefill. Cases cover cache merging and optional attention semantics in subgroup and portable forms.", + "priority": 34, + "when": ["sharedKvFloatOk", "prefillFlashShapeOk", "cachedClusterStorageOk if \"\" == \"\" else cachedNoSgClusterStorageOk", "not present.cosCacheT and not rotaryRequested", "not present.qNormWeightT", "not present.attentionBiasT", "present.headSinkT", "standardSoftmax", "(subgroupCachedPrefillCluster and cachedF32ClusterRegisterGeometry) if \"\" == \"\" else true"], + "constants": { + "scalar": "gqaScalar", + "usesF16": "tensorDtypes.queryT == \"float16\"", + "hasCausal": true, + "causalRightAlign": true, + "hasBias": false, + "hasHeadSink": true, + "hasMask": false, + "maskIsBool": false, + "hasWindow": "attrs.local_window_size > 0", + "headDim": "gqaHeadDim", + "headDimV4": "gqaHeadDim / 4", + "qHiddenV4": "dim(shapes.queryT, 2) / 4", + "qNumHeads": "attrs.num_heads", + "kvNumHeads": "attrs.kv_num_heads", + "TILE_Q": "cachedPrefillQueryTile", + "TILE_K": "cachedClusterTileK", + "LPQ": "cachedPrefillLanesPerQuery", + "kvHeads": "attrs.kv_num_heads", + "packed": "gqaHeadDim", + "mode": "\"copy\"", + "inputVec4": "gqaInputVec4", + "inputScalar": "gqaScalar", + "batchNoSgReduction": "\"\" == \"_nosg\"", + "useSubgroups": "\"\" == \"\"" + }, + "passes": [ + { + "id": "present", + "name": "GroupQueryAttention.Present", + "shader": "gqa-present.wgsl.jinja", + "bindings": "presentCopy", + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", + "workgroupSize": "copyWorkgroupSize" + }, + "viewAlias": [{ "input": "src_k", "output": "present_key" }, { "input": "src_v", "output": "present_value" }] + }, + { + "id": "attention", + "name": "GroupQueryAttention.FlashPrefill", + "source": { + "shader": "attn-flash-prefill-cluster.wgsl.jinja", + "inputs": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" } + }, + "bindings": "flashPrefillPastHeadSink", + "dispatch": { + "x": "ceilDiv(dim(shapes.queryT, 1), constants.TILE_Q)", + "y": "attrs.num_heads", + "z": "dim(shapes.queryT, 0)" + } + } + ] + }, + { + "id": "past_kv_bias_headsink_flash_prefill", + "description": "Updates the float cache in the requested mode and applies clustered online-softmax causal prefill. Cases cover cache merging and optional attention semantics in subgroup and portable forms.", + "priority": 34, + "when": ["sharedKvFloatOk", "prefillFlashShapeOk", "cachedMaskClusterStorageOk if \"\" == \"\" else cachedNoSgClusterStorageOk", "not present.cosCacheT and not rotaryRequested", "not present.qNormWeightT", "present.attentionBiasT", "present.headSinkT", "standardSoftmax", "(subgroupCachedPrefillCluster and cachedF32ClusterRegisterGeometry) if \"\" == \"\" else true"], + "constants": { + "scalar": "gqaScalar", + "usesF16": "tensorDtypes.queryT == \"float16\"", + "hasCausal": true, + "causalRightAlign": true, + "hasBias": false, + "hasHeadSink": true, + "hasMask": true, + "maskIsBool": false, + "hasWindow": "attrs.local_window_size > 0", + "headDim": "gqaHeadDim", + "headDimV4": "gqaHeadDim / 4", + "qHiddenV4": "dim(shapes.queryT, 2) / 4", + "qNumHeads": "attrs.num_heads", + "kvNumHeads": "attrs.kv_num_heads", + "TILE_Q": "cachedPrefillQueryTile", + "TILE_K": "cachedClusterTileK", + "LPQ": "cachedPrefillLanesPerQuery", + "kvHeads": "attrs.kv_num_heads", + "packed": "gqaHeadDim", + "mode": "\"copy\"", + "inputVec4": "gqaInputVec4", + "inputScalar": "gqaScalar", + "batchNoSgReduction": "\"\" == \"_nosg\"", + "useSubgroups": "\"\" == \"\"" + }, + "passes": [ + { + "id": "present", + "name": "GroupQueryAttention.Present", + "shader": "gqa-present.wgsl.jinja", + "bindings": "presentCopy", + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", + "workgroupSize": "copyWorkgroupSize" + }, + "viewAlias": [{ "input": "src_k", "output": "present_key" }, { "input": "src_v", "output": "present_value" }] + }, + { + "id": "attention", + "name": "GroupQueryAttention.FlashPrefill", + "source": { + "shader": "attn-flash-prefill-cluster.wgsl.jinja", + "inputs": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" } + }, + "bindings": "flashPrefillBiasHeadSink", + "dispatch": { + "x": "ceilDiv(dim(shapes.queryT, 1), constants.TILE_Q)", + "y": "attrs.num_heads", + "z": "dim(shapes.queryT, 0)" + } + } + ] + }, + { + "id": "window_shift_flash_prefill", + "description": "Updates the float cache in the requested mode and applies clustered online-softmax causal prefill. Cases cover cache merging and optional attention semantics in subgroup and portable forms.", + "priority": 34, + "when": ["windowShiftOk", "prefillFlashShapeOk", "cachedClusterStorageOk if \"\" == \"\" else cachedNoSgClusterStorageOk", "plainAttentionOptions", "(subgroupCachedPrefillCluster and cachedF32ClusterRegisterGeometry) if \"\" == \"\" else true"], + "constants": { + "scalar": "gqaScalar", + "usesF16": "tensorDtypes.queryT == \"float16\"", + "hasCausal": true, + "causalRightAlign": true, + "hasBias": false, + "hasMask": false, + "maskIsBool": false, + "hasWindow": "attrs.local_window_size > 0", + "headDim": "gqaHeadDim", + "headDimV4": "gqaHeadDim / 4", + "qHiddenV4": "dim(shapes.queryT, 2) / 4", + "qNumHeads": "attrs.num_heads", + "kvNumHeads": "attrs.kv_num_heads", + "TILE_Q": "cachedPrefillQueryTile", + "TILE_K": "cachedClusterTileK", + "LPQ": "cachedPrefillLanesPerQuery", + "kvHeads": "attrs.kv_num_heads", + "packed": "gqaHeadDim", + "mode": "\"window_shift\"", + "useSeqlens": true, + "zeroScalar": "gqaScalar", + "inputVec4": "gqaInputVec4", + "inputScalar": "gqaScalar", + "batchNoSgReduction": "\"\" == \"_nosg\"", + "useSubgroups": "\"\" == \"\"" + }, + "passes": [ + { + "id": "present", + "name": "GroupQueryAttention.WindowShift", + "shader": "gqa-present.wgsl.jinja", + "bindings": "presentMergeShare", + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)", + "workgroupSize": "copyWorkgroupSize" + } + }, + { + "id": "attention", + "name": "GroupQueryAttention.FlashPrefill", + "source": { + "shader": "attn-flash-prefill-cluster.wgsl.jinja", + "inputs": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" } + }, + "bindings": "flashPrefillSeqlens", + "dispatch": { + "x": "ceilDiv(dim(shapes.queryT, 1), constants.TILE_Q)", + "y": "attrs.num_heads", + "z": "dim(shapes.queryT, 0)" + } + } + ] + }, + { + "id": "past_kv_flash_prefill_nosg", + "description": "Updates the float cache in the requested mode and applies clustered online-softmax causal prefill. Cases cover cache merging and optional attention semantics in subgroup and portable forms.", + "priority": 34, + "when": ["sharedKvFloatOk", "prefillFlashShapeOk", "cachedClusterStorageOk if \"_nosg\" == \"\" else cachedNoSgClusterStorageOk", "plainAttentionOptions", "(subgroupCachedPrefillCluster and cachedF32ClusterRegisterGeometry) if \"_nosg\" == \"\" else true"], + "constants": { + "scalar": "gqaScalar", + "usesF16": "tensorDtypes.queryT == \"float16\"", + "hasCausal": true, + "causalRightAlign": true, + "hasBias": false, + "hasMask": false, + "maskIsBool": false, + "hasWindow": "attrs.local_window_size > 0", + "headDim": "gqaHeadDim", + "headDimV4": "gqaHeadDim / 4", + "qHiddenV4": "dim(shapes.queryT, 2) / 4", + "qNumHeads": "attrs.num_heads", + "kvNumHeads": "attrs.kv_num_heads", + "TILE_Q": "cachedPrefillQueryTile", + "TILE_K": "cachedNoSgTileK", + "LPQ": "cachedPrefillLanesPerQuery", + "kvHeads": "attrs.kv_num_heads", + "packed": "gqaHeadDim", + "mode": "\"copy\"", + "inputVec4": "gqaInputVec4", + "inputScalar": "gqaScalar", + "batchNoSgReduction": "\"_nosg\" == \"_nosg\"", + "useSubgroups": "\"_nosg\" == \"\"" + }, + "passes": [ + { + "id": "present", + "name": "GroupQueryAttention.Present", + "shader": "gqa-present.wgsl.jinja", + "bindings": "presentCopy", + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", + "workgroupSize": "copyWorkgroupSize" + }, + "viewAlias": [{ "input": "src_k", "output": "present_key" }, { "input": "src_v", "output": "present_value" }] + }, + { + "id": "attention", + "name": "GroupQueryAttention.FlashPrefill", + "source": { + "shader": "attn-flash-prefill-cluster.wgsl.jinja", + "inputs": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" } + }, + "bindings": "flashPrefillPast", + "dispatch": { + "x": "ceilDiv(dim(shapes.queryT, 1), constants.TILE_Q)", + "y": "attrs.num_heads", + "z": "dim(shapes.queryT, 0)" + } + } + ] + }, + { + "id": "new_kv_past_flash_prefill_nosg", + "description": "Updates the float cache in the requested mode and applies clustered online-softmax causal prefill. Cases cover cache merging and optional attention semantics in subgroup and portable forms.", + "priority": 34, + "when": ["newKvPastOk", "prefillFlashShapeOk", "cachedClusterStorageOk if \"_nosg\" == \"\" else cachedNoSgClusterStorageOk", "plainAttentionOptions", "(subgroupCachedPrefillCluster and cachedF32ClusterRegisterGeometry) if \"_nosg\" == \"\" else true"], + "constants": { + "scalar": "gqaScalar", + "usesF16": "tensorDtypes.queryT == \"float16\"", + "hasCausal": true, + "causalRightAlign": true, + "hasBias": false, + "hasMask": false, + "maskIsBool": false, + "hasWindow": "attrs.local_window_size > 0", + "headDim": "gqaHeadDim", + "headDimV4": "gqaHeadDim / 4", + "qHiddenV4": "dim(shapes.queryT, 2) / 4", + "qNumHeads": "attrs.num_heads", + "kvNumHeads": "attrs.kv_num_heads", + "TILE_Q": "cachedPrefillQueryTile", + "TILE_K": "cachedNoSgTileK", + "LPQ": "cachedPrefillLanesPerQuery", + "kvHeads": "attrs.kv_num_heads", + "packed": "gqaHeadDim", + "mode": "\"merge\"", + "inputVec4": "gqaInputVec4", + "inputScalar": "gqaScalar", + "batchNoSgReduction": "\"_nosg\" == \"_nosg\"", + "useSubgroups": "\"_nosg\" == \"\"" + }, + "passes": [ + { + "id": "present", + "name": "GroupQueryAttention.Merge", + "shader": "gqa-present.wgsl.jinja", + "bindings": "presentMerge", + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)", + "workgroupSize": "copyWorkgroupSize" + } + }, + { + "id": "attention", + "name": "GroupQueryAttention.FlashPrefill", + "source": { + "shader": "attn-flash-prefill-cluster.wgsl.jinja", + "inputs": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" } + }, + "bindings": "flashPrefillNew", + "dispatch": { + "x": "ceilDiv(dim(shapes.queryT, 1), constants.TILE_Q)", + "y": "attrs.num_heads", + "z": "dim(shapes.queryT, 0)" + } + } + ] + }, + { + "id": "past_kv_rotary_flash_prefill_nosg", + "description": "Updates the float cache in the requested mode and applies clustered online-softmax causal prefill. Cases cover cache merging and optional attention semantics in subgroup and portable forms.", + "priority": 36, + "when": ["sharedKvFloatOk", "prefillFlashShapeOk", "cachedClusterStorageOk if \"_nosg\" == \"\" else cachedNoSgClusterStorageOk", "standardSoftmax", "present.cosCacheT and rotaryRequested", "not present.qNormWeightT", "not present.attentionBiasT", "not present.headSinkT", "(subgroupCachedPrefillCluster and cachedF32ClusterRegisterGeometry) if \"_nosg\" == \"\" else true"], + "constants": { + "scalar": "gqaScalar", + "usesF16": "tensorDtypes.queryT == \"float16\"", + "hasCausal": true, + "causalRightAlign": true, + "hasBias": false, + "hasMask": false, + "maskIsBool": false, + "hasWindow": "attrs.local_window_size > 0", + "headDim": "gqaHeadDim", + "headDimV4": "gqaHeadDim / 4", + "qHiddenV4": "dim(shapes.queryT, 2) / 4", + "qNumHeads": "attrs.num_heads", + "kvNumHeads": "attrs.kv_num_heads", + "TILE_Q": "cachedPrefillQueryTile", + "TILE_K": "cachedNoSgTileK", + "LPQ": "cachedPrefillLanesPerQuery", + "kvHeads": "attrs.kv_num_heads", + "packed": "gqaHeadDim", + "mode": "\"copy\"", + "hasRotary": true, + "fusedRotary": true, + "half": "gqaHeadDim / 2", + "qHeads": "attrs.num_heads", + "qHidden": "dim(shapes.queryT, 2)", + "inputVec4": "gqaInputVec4", + "inputScalar": "gqaScalar", + "cosScalar": "\"f16\" if tensorDtypes.cosCacheT == \"float16\" else \"f32\"", + "batchNoSgReduction": "\"_nosg\" == \"_nosg\"", + "useSubgroups": "\"_nosg\" == \"\"" + }, + "passes": [ + { + "id": "present", + "name": "GroupQueryAttention.Present", + "shader": "gqa-present.wgsl.jinja", + "bindings": "presentCopy", + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", + "workgroupSize": "copyWorkgroupSize" + }, + "viewAlias": [{ "input": "src_k", "output": "present_key" }, { "input": "src_v", "output": "present_value" }] + }, + { + "id": "attention", + "name": "GroupQueryAttention.FlashPrefill", + "source": { + "shader": "attn-flash-prefill-cluster.wgsl.jinja", + "inputs": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" } + }, + "bindings": "rotaryFusedFlashAttention", + "dispatch": { + "x": "ceilDiv(dim(shapes.queryT, 1), constants.TILE_Q)", + "y": "attrs.num_heads", + "z": "dim(shapes.queryT, 0)" + } + } + ] + }, + { + "id": "past_kv_softcap_flash_prefill_nosg", + "description": "Updates the float cache in the requested mode and applies clustered online-softmax causal prefill. Cases cover cache merging and optional attention semantics in subgroup and portable forms.", + "priority": 34, + "when": ["sharedKvFloatOk", "prefillFlashShapeOk", "cachedClusterStorageOk if \"_nosg\" == \"\" else cachedNoSgClusterStorageOk", "noAuxAttentionInputs", "attrs.softcap > 0", "attrs.smooth_softmax != 1", "(subgroupCachedPrefillCluster and cachedF32ClusterRegisterGeometry) if \"_nosg\" == \"\" else true"], + "constants": { + "scalar": "gqaScalar", + "usesF16": "tensorDtypes.queryT == \"float16\"", + "hasCausal": true, + "causalRightAlign": true, + "hasBias": false, + "hasSoftcap": true, + "hasMask": false, + "maskIsBool": false, + "hasWindow": "attrs.local_window_size > 0", + "headDim": "gqaHeadDim", + "headDimV4": "gqaHeadDim / 4", + "qHiddenV4": "dim(shapes.queryT, 2) / 4", + "qNumHeads": "attrs.num_heads", + "kvNumHeads": "attrs.kv_num_heads", + "TILE_Q": "cachedPrefillQueryTile", + "TILE_K": "cachedNoSgTileK", + "LPQ": "cachedPrefillLanesPerQuery", + "kvHeads": "attrs.kv_num_heads", + "packed": "gqaHeadDim", + "mode": "\"copy\"", + "inputVec4": "gqaInputVec4", + "inputScalar": "gqaScalar", + "batchNoSgReduction": "\"_nosg\" == \"_nosg\"", + "useSubgroups": "\"_nosg\" == \"\"" + }, + "passes": [ + { + "id": "present", + "name": "GroupQueryAttention.Present", + "shader": "gqa-present.wgsl.jinja", + "bindings": "presentCopy", + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", + "workgroupSize": "copyWorkgroupSize" + }, + "viewAlias": [{ "input": "src_k", "output": "present_key" }, { "input": "src_v", "output": "present_value" }] + }, + { + "id": "attention", + "name": "GroupQueryAttention.FlashPrefill", + "source": { + "shader": "attn-flash-prefill-cluster.wgsl.jinja", + "inputs": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" } + }, + "bindings": "flashPrefillPastSoftcap", + "dispatch": { + "x": "ceilDiv(dim(shapes.queryT, 1), constants.TILE_Q)", + "y": "attrs.num_heads", + "z": "dim(shapes.queryT, 0)" + } + } + ] + }, + { + "id": "past_kv_headsink_flash_prefill_nosg", + "description": "Updates the float cache in the requested mode and applies clustered online-softmax causal prefill. Cases cover cache merging and optional attention semantics in subgroup and portable forms.", + "priority": 34, + "when": ["sharedKvFloatOk", "prefillFlashShapeOk", "cachedClusterStorageOk if \"_nosg\" == \"\" else cachedNoSgClusterStorageOk", "not present.cosCacheT and not rotaryRequested", "not present.qNormWeightT", "not present.attentionBiasT", "present.headSinkT", "standardSoftmax", "(subgroupCachedPrefillCluster and cachedF32ClusterRegisterGeometry) if \"_nosg\" == \"\" else true"], + "constants": { + "scalar": "gqaScalar", + "usesF16": "tensorDtypes.queryT == \"float16\"", + "hasCausal": true, + "causalRightAlign": true, + "hasBias": false, + "hasHeadSink": true, + "hasMask": false, + "maskIsBool": false, + "hasWindow": "attrs.local_window_size > 0", + "headDim": "gqaHeadDim", + "headDimV4": "gqaHeadDim / 4", + "qHiddenV4": "dim(shapes.queryT, 2) / 4", + "qNumHeads": "attrs.num_heads", + "kvNumHeads": "attrs.kv_num_heads", + "TILE_Q": "cachedPrefillQueryTile", + "TILE_K": "cachedNoSgTileK", + "LPQ": "cachedPrefillLanesPerQuery", + "kvHeads": "attrs.kv_num_heads", + "packed": "gqaHeadDim", + "mode": "\"copy\"", + "inputVec4": "gqaInputVec4", + "inputScalar": "gqaScalar", + "batchNoSgReduction": "\"_nosg\" == \"_nosg\"", + "useSubgroups": "\"_nosg\" == \"\"" + }, + "passes": [ + { + "id": "present", + "name": "GroupQueryAttention.Present", + "shader": "gqa-present.wgsl.jinja", + "bindings": "presentCopy", + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", + "workgroupSize": "copyWorkgroupSize" + }, + "viewAlias": [{ "input": "src_k", "output": "present_key" }, { "input": "src_v", "output": "present_value" }] + }, + { + "id": "attention", + "name": "GroupQueryAttention.FlashPrefill", + "source": { + "shader": "attn-flash-prefill-cluster.wgsl.jinja", + "inputs": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" } + }, + "bindings": "flashPrefillPastHeadSink", + "dispatch": { + "x": "ceilDiv(dim(shapes.queryT, 1), constants.TILE_Q)", + "y": "attrs.num_heads", + "z": "dim(shapes.queryT, 0)" + } + } + ] + }, + { + "id": "past_kv_bias_headsink_flash_prefill_nosg", + "description": "Updates the float cache in the requested mode and applies clustered online-softmax causal prefill. Cases cover cache merging and optional attention semantics in subgroup and portable forms.", + "priority": 34, + "when": ["sharedKvFloatOk", "prefillFlashShapeOk", "cachedMaskClusterStorageOk if \"_nosg\" == \"\" else cachedNoSgClusterStorageOk", "not present.cosCacheT and not rotaryRequested", "not present.qNormWeightT", "present.attentionBiasT", "present.headSinkT", "standardSoftmax", "(subgroupCachedPrefillCluster and cachedF32ClusterRegisterGeometry) if \"_nosg\" == \"\" else true"], + "constants": { + "scalar": "gqaScalar", + "usesF16": "tensorDtypes.queryT == \"float16\"", + "hasCausal": true, + "causalRightAlign": true, + "hasBias": false, + "hasHeadSink": true, + "hasMask": true, + "maskIsBool": false, + "hasWindow": "attrs.local_window_size > 0", + "headDim": "gqaHeadDim", + "headDimV4": "gqaHeadDim / 4", + "qHiddenV4": "dim(shapes.queryT, 2) / 4", + "qNumHeads": "attrs.num_heads", + "kvNumHeads": "attrs.kv_num_heads", + "TILE_Q": "cachedPrefillQueryTile", + "TILE_K": "cachedNoSgTileK", + "LPQ": "cachedPrefillLanesPerQuery", + "kvHeads": "attrs.kv_num_heads", + "packed": "gqaHeadDim", + "mode": "\"copy\"", + "inputVec4": "gqaInputVec4", + "inputScalar": "gqaScalar", + "batchNoSgReduction": "\"_nosg\" == \"_nosg\"", + "useSubgroups": "\"_nosg\" == \"\"" + }, + "passes": [ + { + "id": "present", + "name": "GroupQueryAttention.Present", + "shader": "gqa-present.wgsl.jinja", + "bindings": "presentCopy", + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", + "workgroupSize": "copyWorkgroupSize" + }, + "viewAlias": [{ "input": "src_k", "output": "present_key" }, { "input": "src_v", "output": "present_value" }] + }, + { + "id": "attention", + "name": "GroupQueryAttention.FlashPrefill", + "source": { + "shader": "attn-flash-prefill-cluster.wgsl.jinja", + "inputs": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" } + }, + "bindings": "flashPrefillBiasHeadSink", + "dispatch": { + "x": "ceilDiv(dim(shapes.queryT, 1), constants.TILE_Q)", + "y": "attrs.num_heads", + "z": "dim(shapes.queryT, 0)" + } + } + ] + }, + { + "id": "window_shift_flash_prefill_nosg", + "description": "Updates the float cache in the requested mode and applies clustered online-softmax causal prefill. Cases cover cache merging and optional attention semantics in subgroup and portable forms.", + "priority": 34, + "when": ["windowShiftOk", "prefillFlashShapeOk", "cachedClusterStorageOk if \"_nosg\" == \"\" else cachedNoSgClusterStorageOk", "plainAttentionOptions", "(subgroupCachedPrefillCluster and cachedF32ClusterRegisterGeometry) if \"_nosg\" == \"\" else true"], + "constants": { + "scalar": "gqaScalar", + "usesF16": "tensorDtypes.queryT == \"float16\"", + "hasCausal": true, + "causalRightAlign": true, + "hasBias": false, + "hasMask": false, + "maskIsBool": false, + "hasWindow": "attrs.local_window_size > 0", + "headDim": "gqaHeadDim", + "headDimV4": "gqaHeadDim / 4", + "qHiddenV4": "dim(shapes.queryT, 2) / 4", + "qNumHeads": "attrs.num_heads", + "kvNumHeads": "attrs.kv_num_heads", + "TILE_Q": "cachedPrefillQueryTile", + "TILE_K": "cachedNoSgTileK", + "LPQ": "cachedPrefillLanesPerQuery", + "kvHeads": "attrs.kv_num_heads", + "packed": "gqaHeadDim", + "mode": "\"window_shift\"", + "useSeqlens": true, + "zeroScalar": "gqaScalar", + "inputVec4": "gqaInputVec4", + "inputScalar": "gqaScalar", + "batchNoSgReduction": "\"_nosg\" == \"_nosg\"", + "useSubgroups": "\"_nosg\" == \"\"" + }, + "passes": [ + { + "id": "present", + "name": "GroupQueryAttention.WindowShift", + "shader": "gqa-present.wgsl.jinja", + "bindings": "presentMergeShare", + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)", + "workgroupSize": "copyWorkgroupSize" + } + }, + { + "id": "attention", + "name": "GroupQueryAttention.FlashPrefill", + "source": { + "shader": "attn-flash-prefill-cluster.wgsl.jinja", + "inputs": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" } + }, + "bindings": "flashPrefillSeqlens", + "dispatch": { + "x": "ceilDiv(dim(shapes.queryT, 1), constants.TILE_Q)", + "y": "attrs.num_heads", + "z": "dim(shapes.queryT, 0)" + } + } + ] + }, + { + "id": "qkv_present_flash_q32_broadcast", + "priority": 30, + "requires": { "features": ["subgroups"] }, + "when": ["subgroupsWave32", "device.features.has(\"shader-f16\")", "tensorDtypes.query == \"float16\"", "flashShapeOk", "headDim % 32 == 0", "headDim >= 64", "headDim <= 256", "dim(shapes.query, 1) >= tunables.QKV_PREFILL_MIN_QUERY_TOKENS", "attrs.local_window_size <= 0"], + "constants": { + "hasBias": false, + "hasCausal": false, + "usesF16": true, + "scalar": "\"f16\"", + "inputVec4": "\"vec4\"", + "outputScalar": "\"f16\"", + "presentScalar": "qkvCacheScalar", + "inputElement": "\"vec4\"", + "outputElement": "\"vec4\"", + "presentElement": "qkvCacheVec4", + "qNumHeads": "attrs.num_heads", + "kvNumHeads": "attrs.kv_num_heads", + "headDim": "dim(shapes.query, 2) / attrs.num_heads", + "headDimV4": "(dim(shapes.query, 2) / attrs.num_heads) / 4", + "kvHidden": "dim(shapes.key, 2)", + "kvHiddenV4": "dim(shapes.key, 2) / 4", + "copyWorkgroupSize": "copyWorkgroupSize", + "presentVec4": true, + "presentElem": "qkvCacheVec4", + "kStep": 64, + "qkGroups": 16 + }, + "passes": [ + { + "id": "attention", + "name": "GroupQueryAttention.FlashQ32Broadcast", + "source": { "shader": "attn-flash-q32-broadcast.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, + "constants": { "usesF16": "tensorDtypes.query == \"float16\"" }, + "bindings": "qkvAttentionNoWindow", + "dispatch": { "x": "ceilDiv(dim(shapes.query, 1), 32)", "y": "attrs.num_heads", "z": "dim(shapes.query, 0)" } + }, + { + "id": "present", + "name": "GroupQueryAttention.Present", + "shader": "gqa-present.wgsl.jinja", + "bindings": "qkvPresent", + "dispatch": { "threads": "numel(shapes.presentKeyT) / 4", "workgroupSize": "constants.copyWorkgroupSize" } + } + ] + }, + { + "id": "qkv_present_flash_q32_shared", + "priority": 29, + "requires": { "features": ["shader-f16"] }, + "when": ["tensorDtypes.query == \"float16\"", "flashShapeOk", "headDim % 4 == 0", "headDim >= 64", "headDim <= 256", "dim(shapes.query, 1) >= tunables.QKV_PREFILL_MIN_QUERY_TOKENS", "attrs.local_window_size <= 0", "ceilDiv(dim(shapes.query, 1), 64) <= device.limits.maxComputeWorkgroupsPerDimension", "(headDim / 4) * 32 * 16 <= device.limits.maxComputeWorkgroupStorageSize"], + "constants": { + "hasBias": false, + "hasCausal": false, + "usesF16": true, + "scalar": "\"f16\"", + "inputVec4": "\"vec4\"", + "outputScalar": "\"f16\"", + "presentScalar": "qkvCacheScalar", + "inputElement": "\"vec4\"", + "outputElement": "\"vec4\"", + "presentElement": "qkvCacheVec4", + "qNumHeads": "attrs.num_heads", + "kvNumHeads": "attrs.kv_num_heads", + "headDim": "dim(shapes.query, 2) / attrs.num_heads", + "headDimV4": "(dim(shapes.query, 2) / attrs.num_heads) / 4", + "kvHidden": "dim(shapes.key, 2)", + "kvHiddenV4": "dim(shapes.key, 2) / 4", + "copyWorkgroupSize": "copyWorkgroupSize", + "presentVec4": true, + "presentElem": "qkvCacheVec4", + "kStep": 32, + "qkGroups": 8, + "qStep": 64 + }, + "passes": [ + { + "id": "attention", + "name": "GroupQueryAttention.FlashQ32Shared", + "source": { + "shader": "attn-flash-q32-broadcast.wgsl.jinja", + "inputs": { "layout": "\"bsh\"", "useSubgroups": "false" } + }, + "constants": { "usesF16": "tensorDtypes.query == \"float16\"" }, + "bindings": "qkvAttentionNoWindow", + "dispatch": { "x": "ceilDiv(dim(shapes.query, 1), 64)", "y": "attrs.num_heads", "z": "dim(shapes.query, 0)" } + }, + { + "id": "present", + "name": "GroupQueryAttention.Present", + "shader": "gqa-present.wgsl.jinja", + "bindings": "qkvPresent", + "dispatch": { "threads": "numel(shapes.presentKeyT) / 4", "workgroupSize": "constants.copyWorkgroupSize" } + } + ] + }, + { + "id": "past_kv", + "priority": 30, + "when": ["sharedKvFloatOk", "gqaScalarDispatchOk", "noAuxAttentionInputs", "gqaHeadDim % 2 == 0"], + "constants": { + "useSeqlens": true, + "headDim": "gqaHeadDim", + "cooperative": "gqaCoop", + "dispatchUnits": "gqaDispatchUnits", + "qHeads": "attrs.num_heads", + "kvHeads": "attrs.kv_num_heads", + "qHidden": "dim(shapes.queryT, 2)", + "packed": "gqaHeadDim", + "mode": "\"copy\"", + "usesF16": "tensorDtypes.queryT == \"float16\"", + "inputScalar": "gqaScalar" + }, + "passes": [ + { + "id": "attention", + "name": "GroupQueryAttention.Attn", + "shader": "gqa-attention.wgsl.jinja", + "bindings": "cachedAttention", + "dispatch": { "workgroups": "constants.dispatchUnits" } + }, + { + "id": "present", + "name": "GroupQueryAttention.Present", + "shader": "gqa-present.wgsl.jinja", + "bindings": "presentCopy", + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", + "workgroupSize": "copyWorkgroupSize" + }, + "viewAlias": [{ "input": "src_k", "output": "present_key" }, { "input": "src_v", "output": "present_value" }] + } + ] + }, + { + "id": "past_kv_rotary", + "priority": 31, + "when": ["sharedKvFloatOk", "gqaScalarDispatchOk", "gqaHeadDim % 2 == 0", "present.cosCacheT and rotaryRequested", "not present.qNormWeightT"], + "constants": { + "useSeqlens": true, + "headDim": "gqaHeadDim", + "cooperative": "gqaCoop", + "dispatchUnits": "gqaDispatchUnits", + "qHeads": "attrs.num_heads", + "kvHeads": "attrs.kv_num_heads", + "qHidden": "dim(shapes.queryT, 2)", + "packed": "gqaHeadDim", + "half": "gqaHeadDim / 2", + "hasRotary": true, + "mode": "\"copy\"", + "usesF16": "tensorDtypes.queryT == \"float16\"", + "inputScalar": "gqaScalar", + "cosScalar": "\"f16\" if tensorDtypes.cosCacheT == \"float16\" else \"f32\"" + }, + "passes": [ + { + "id": "attention", + "name": "GroupQueryAttention.Attn", + "shader": "gqa-attention.wgsl.jinja", + "bindings": "cachedRotaryAttention", + "dispatch": { "workgroups": "constants.dispatchUnits" } + }, + { + "id": "present", + "name": "GroupQueryAttention.Present", + "shader": "gqa-present.wgsl.jinja", + "bindings": "presentCopy", + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", + "workgroupSize": "copyWorkgroupSize" + }, + "viewAlias": [{ "input": "src_k", "output": "present_key" }, { "input": "src_v", "output": "present_value" }] + } + ] + }, + { + "id": "past_kv_qnorm_rotary", + "priority": 32, + "when": ["sharedKvFloatOk", "gqaScalarDispatchOk", "gqaHeadDim % 2 == 0", "present.cosCacheT and rotaryRequested", "present.qNormWeightT"], + "constants": { + "useSeqlens": true, + "headDim": "gqaHeadDim", + "cooperative": "gqaCoop", + "dispatchUnits": "gqaDispatchUnits", + "qHeads": "attrs.num_heads", + "kvHeads": "attrs.kv_num_heads", + "qHidden": "dim(shapes.queryT, 2)", + "packed": "gqaHeadDim", + "half": "gqaHeadDim / 2", + "qkEps": "attrs.qk_norm_epsilon", + "hasRotary": true, + "hasQNorm": true, + "mode": "\"copy\"", + "usesF16": "tensorDtypes.queryT == \"float16\"", + "inputScalar": "gqaScalar", + "cosScalar": "\"f16\" if tensorDtypes.cosCacheT == \"float16\" else \"f32\"", + "qNormScalar": "\"f16\" if tensorDtypes.qNormWeightT == \"float16\" else \"f32\"" + }, + "passes": [ + { + "id": "attention", + "name": "GroupQueryAttention.Attn", + "shader": "gqa-attention.wgsl.jinja", + "bindings": "cachedQnormRotaryAttention", + "dispatch": { "workgroups": "constants.dispatchUnits" } + }, + { + "id": "present", + "name": "GroupQueryAttention.Present", + "shader": "gqa-present.wgsl.jinja", + "bindings": "presentCopy", + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", + "workgroupSize": "copyWorkgroupSize" + }, + "viewAlias": [{ "input": "src_k", "output": "present_key" }, { "input": "src_v", "output": "present_value" }] + } + ] + }, + { + "id": "new_kv_past", + "priority": 30, + "when": ["newKvPastOk", "gqaScalarDispatchOk", "noAuxAttentionInputs", "gqaHeadDim % 2 == 0"], + "constants": { + "headDim": "gqaHeadDim", + "qHeads": "attrs.num_heads", + "kvHeads": "attrs.kv_num_heads", + "qHidden": "dim(shapes.queryT, 2)", + "packed": "gqaHeadDim", + "mode": "\"merge\"", + "cooperative": "gqaCoop", + "dispatchUnits": "gqaDispatchUnits", + "usesF16": "tensorDtypes.queryT == \"float16\"", + "inputScalar": "gqaScalar" + }, + "passes": [ + { + "id": "present", + "name": "GroupQueryAttention.Merge", + "shader": "gqa-present.wgsl.jinja", + "bindings": "presentMerge", + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)", + "workgroupSize": "copyWorkgroupSize" + } + }, + { + "id": "attention", + "name": "GroupQueryAttention.Attn", + "shader": "gqa-attention.wgsl.jinja", + "bindings": "mergedAttention", + "dispatch": { "workgroups": "constants.dispatchUnits" } + } + ] + }, + { + "id": "window_shift_append", + "priority": 30, + "when": ["windowShiftOk", "gqaScalarDispatchOk", "noAuxAttentionInputs", "gqaHeadDim % 2 == 0"], + "constants": { + "useSeqlens": true, + "headDim": "gqaHeadDim", + "cooperative": "gqaCoop", + "dispatchUnits": "gqaDispatchUnits", + "qHeads": "attrs.num_heads", + "kvHeads": "attrs.kv_num_heads", + "qHidden": "dim(shapes.queryT, 2)", + "packed": "gqaHeadDim", + "mode": "\"window_shift\"", + "usesF16": "tensorDtypes.queryT == \"float16\"", + "inputScalar": "gqaScalar", + "zeroScalar": "gqaScalar" + }, + "passes": [ + { + "id": "present", + "name": "GroupQueryAttention.WindowShift", + "shader": "gqa-present.wgsl.jinja", + "bindings": "presentMergeShare", + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)", + "workgroupSize": "copyWorkgroupSize" + } + }, + { + "id": "attention", + "name": "GroupQueryAttention.Attn", + "shader": "gqa-attention.wgsl.jinja", + "bindings": "mergedSeqlensAttention", + "dispatch": { "workgroups": "constants.dispatchUnits" } + } + ] + }, + { + "id": "new_kv_qnorm_rotary", + "priority": 33, + "when": ["copyWorkgroupOk", "gqaScalarDispatchOk", "present.seqlensKT", "not present.pastKeyT", "dim(shapes.keyT, 1) > 0", "tensorDtypes.queryT == \"float32\"", "not present.kScaleT", "present.cosCacheT and rotaryRequested", "present.qNormWeightT", "dim(shapes.queryT, 2) % attrs.num_heads == 0", "gqaHeadDim % 2 == 0"], + "constants": { + "headDim": "gqaHeadDim", + "cooperative": "gqaCoop", + "dispatchUnits": "gqaDispatchUnits", + "qHeads": "attrs.num_heads", + "kvHeads": "attrs.kv_num_heads", + "qHidden": "dim(shapes.queryT, 2)", + "packed": "gqaHeadDim", + "half": "gqaHeadDim / 2", + "qkEps": "attrs.qk_norm_epsilon", + "hasRotary": true, + "hasQNorm": true, + "mode": "\"build\"", + "hasKNorm": true + }, + "passes": [ + { + "id": "present", + "name": "GroupQueryAttention.Present", + "shader": "gqa-present.wgsl.jinja", + "bindings": "newQnormRotaryPresent", + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)", + "workgroupSize": "copyWorkgroupSize" + } + }, + { + "id": "attention", + "name": "GroupQueryAttention.Attn", + "shader": "gqa-attention.wgsl.jinja", + "bindings": "newQnormRotaryAttention", + "dispatch": { "workgroups": "constants.dispatchUnits" } + } + ] + }, + { + "id": "past_kv_bias", + "priority": 30, + "when": ["sharedKvCacheOk", "gqaScalarDispatchOk", "tensorDtypes.queryT == \"float32\"", "gqaHeadDim % 2 == 0", "not present.cosCacheT and not rotaryRequested", "not present.qNormWeightT", "present.attentionBiasT", "not present.headSinkT"], + "constants": { + "useSeqlens": true, + "headDim": "gqaHeadDim", + "cooperative": "gqaCoop", + "dispatchUnits": "gqaDispatchUnits", + "qHeads": "attrs.num_heads", + "kvHeads": "attrs.kv_num_heads", + "qHidden": "dim(shapes.queryT, 2)", + "packed": "gqaHeadDim", + "hasBias": true, + "mode": "\"copy\"", + "inputScalar": "gqaScalar" + }, + "passes": [ + { + "id": "attention", + "name": "GroupQueryAttention.Attn", + "shader": "gqa-attention.wgsl.jinja", + "bindings": "cachedBiasAttention", + "dispatch": { "workgroups": "constants.dispatchUnits" } + }, + { + "id": "present", + "name": "GroupQueryAttention.Present", + "shader": "gqa-present.wgsl.jinja", + "bindings": "presentCopy", + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", + "workgroupSize": "copyWorkgroupSize" + }, + "viewAlias": [{ "input": "src_k", "output": "present_key" }, { "input": "src_v", "output": "present_value" }] + } + ] + }, + { + "id": "past_kv_headsink", + "priority": 30, + "when": ["sharedKvCacheOk", "gqaScalarDispatchOk", "tensorDtypes.queryT == \"float32\"", "gqaHeadDim % 2 == 0", "not present.cosCacheT and not rotaryRequested", "not present.qNormWeightT", "not present.attentionBiasT", "present.headSinkT"], + "constants": { + "useSeqlens": true, + "headDim": "gqaHeadDim", + "cooperative": "gqaCoop", + "dispatchUnits": "gqaDispatchUnits", + "qHeads": "attrs.num_heads", + "kvHeads": "attrs.kv_num_heads", + "qHidden": "dim(shapes.queryT, 2)", + "packed": "gqaHeadDim", + "hasHeadSink": true, + "mode": "\"copy\"", + "inputScalar": "gqaScalar" + }, + "passes": [ + { + "id": "attention", + "name": "GroupQueryAttention.Attn", + "shader": "gqa-attention.wgsl.jinja", + "bindings": "cachedHeadSinkAttention", + "dispatch": { "workgroups": "constants.dispatchUnits" } + }, + { + "id": "present", + "name": "GroupQueryAttention.Present", + "shader": "gqa-present.wgsl.jinja", + "bindings": "presentCopy", + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", + "workgroupSize": "copyWorkgroupSize" + }, + "viewAlias": [{ "input": "src_k", "output": "present_key" }, { "input": "src_v", "output": "present_value" }] + } + ] + }, + { + "id": "past_kv_bias_headsink", + "priority": 30, + "when": ["sharedKvCacheOk", "gqaScalarDispatchOk", "tensorDtypes.queryT == \"float32\"", "gqaHeadDim % 2 == 0", "not present.cosCacheT and not rotaryRequested", "not present.qNormWeightT", "present.attentionBiasT", "present.headSinkT"], + "constants": { + "useSeqlens": true, + "headDim": "gqaHeadDim", + "cooperative": "gqaCoop", + "dispatchUnits": "gqaDispatchUnits", + "qHeads": "attrs.num_heads", + "kvHeads": "attrs.kv_num_heads", + "qHidden": "dim(shapes.queryT, 2)", + "packed": "gqaHeadDim", + "hasBias": true, + "hasHeadSink": true, + "mode": "\"copy\"", + "inputScalar": "gqaScalar" + }, + "passes": [ + { + "id": "attention", + "name": "GroupQueryAttention.Attn", + "shader": "gqa-attention.wgsl.jinja", + "bindings": "cachedBiasHeadSinkAttention", + "dispatch": { "workgroups": "constants.dispatchUnits" } + }, + { + "id": "present", + "name": "GroupQueryAttention.Present", + "shader": "gqa-present.wgsl.jinja", + "bindings": "presentCopy", + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", + "workgroupSize": "copyWorkgroupSize" + }, + "viewAlias": [{ "input": "src_k", "output": "present_key" }, { "input": "src_v", "output": "present_value" }] + } + ] + }, + { + "id": "quant_int8_flash_prefill", + "description": "Builds an int8 or int4 cache and applies clustered online-softmax prefill directly from the packed values. The family supplies subgroup and portable reductions.", + "priority": 37, + "when": ["quantizedPromptOk", "prefillFlashShapeOk", "cachedClusterStorageOk if \"\" == \"\" else cachedNoSgClusterStorageOk", "attrs.kv_cache_bit_width == 8", "plainAttentionOptions", "(subgroupCachedPrefillCluster and cachedF32ClusterRegisterGeometry) if \"\" == \"\" else true"], + "constants": { + "scalar": "gqaScalar", + "usesF16": false, + "hasCausal": true, + "causalRightAlign": true, + "hasBias": false, + "hasMask": false, + "maskIsBool": false, + "hasWindow": "attrs.local_window_size > 0", + "headDim": "gqaHeadDim", + "headDimV4": "gqaHeadDim / 4", + "qHiddenV4": "dim(shapes.queryT, 2) / 4", + "qNumHeads": "attrs.num_heads", + "kvNumHeads": "attrs.kv_num_heads", + "TILE_Q": "cachedPrefillQueryTile", + "TILE_K": "cachedClusterTileK", + "LPQ": "cachedPrefillLanesPerQuery", + "kvHeads": "attrs.kv_num_heads", + "inputVec4": "gqaInputVec4", + "inputScalar": "gqaScalar", + "batchNoSgReduction": "\"\" == \"_nosg\"", + "useSubgroups": "\"\" == \"\"", + "quantCacheFormat": "\"int8\"", + "quantCacheElement": "\"i32\"", + "quantCacheReadElement": "\"vec4\"", + "packed": "gqaHeadDim", + "cooperative": false, + "dispatchUnits": 1, + "qHeads": "attrs.num_heads", + "qHidden": "dim(shapes.queryT, 2)", + "mode": "\"build_quant\"", + "bits": 8, + "qmax": 127, + "qmin": -128, + "quantized": true + }, + "passes": [ + { + "id": "present", + "name": "GroupQueryAttention.QuantPresent", + "shader": "gqa-present.wgsl.jinja", + "bindings": "quantPresent", + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)", + "workgroupSize": "copyWorkgroupSize" + } + }, + { + "id": "attention", + "name": "GroupQueryAttention.QuantFlashPrefill", + "source": { + "shader": "attn-flash-prefill-cluster.wgsl.jinja", + "inputs": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" } + }, + "bindings": "flashPrefillQuant", + "dispatch": { + "x": "ceilDiv(dim(shapes.queryT, 1), cachedPrefillQueryTile)", + "y": "attrs.num_heads", + "z": "dim(shapes.queryT, 0)" + } + } + ] + }, + { + "id": "quant_int4_flash_prefill", + "description": "Builds an int8 or int4 cache and applies clustered online-softmax prefill directly from the packed values. The family supplies subgroup and portable reductions.", + "priority": 37, + "when": ["quantizedPromptOk", "prefillFlashShapeOk", "cachedClusterStorageOk if \"\" == \"\" else cachedNoSgClusterStorageOk", "attrs.kv_cache_bit_width == 4", "present.kScaleT", "plainAttentionOptions", "(subgroupCachedPrefillCluster and cachedF32ClusterRegisterGeometry) if \"\" == \"\" else true"], + "constants": { + "scalar": "gqaScalar", + "usesF16": false, + "hasCausal": true, + "causalRightAlign": true, + "hasBias": false, + "hasMask": false, + "maskIsBool": false, + "hasWindow": "attrs.local_window_size > 0", + "headDim": "gqaHeadDim", + "headDimV4": "gqaHeadDim / 4", + "qHiddenV4": "dim(shapes.queryT, 2) / 4", + "qNumHeads": "attrs.num_heads", + "kvNumHeads": "attrs.kv_num_heads", + "TILE_Q": "cachedPrefillQueryTile", + "TILE_K": "cachedClusterTileK", + "LPQ": "cachedPrefillLanesPerQuery", + "kvHeads": "attrs.kv_num_heads", + "inputVec4": "gqaInputVec4", + "inputScalar": "gqaScalar", + "batchNoSgReduction": "\"\" == \"_nosg\"", + "useSubgroups": "\"\" == \"\"", + "quantCacheFormat": "\"int4\"", + "quantCacheElement": "\"u32\"", + "quantCacheReadElement": "\"u32\"", + "packed": "gqaHeadDim / 2", + "cooperative": false, + "dispatchUnits": 1, + "qHeads": "attrs.num_heads", + "qHidden": "dim(shapes.queryT, 2)", + "mode": "\"build_quant\"", + "bits": 4, + "qmax": 7, + "qmin": -8, + "quantized": true + }, + "passes": [ + { + "id": "present", + "name": "GroupQueryAttention.QuantPresent", + "shader": "gqa-present.wgsl.jinja", + "bindings": "quantPresent", + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)", + "workgroupSize": "copyWorkgroupSize" + } + }, + { + "id": "attention", + "name": "GroupQueryAttention.QuantFlashPrefill", + "source": { + "shader": "attn-flash-prefill-cluster.wgsl.jinja", + "inputs": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" } + }, + "bindings": "flashPrefillQuant", + "dispatch": { + "x": "ceilDiv(dim(shapes.queryT, 1), cachedPrefillQueryTile)", + "y": "attrs.num_heads", + "z": "dim(shapes.queryT, 0)" + } + } + ] + }, + { + "id": "quant_int8_flash_prefill_nosg", + "description": "Builds an int8 or int4 cache and applies clustered online-softmax prefill directly from the packed values. The family supplies subgroup and portable reductions.", + "priority": 37, + "when": ["quantizedPromptOk", "prefillFlashShapeOk", "cachedClusterStorageOk if \"_nosg\" == \"\" else cachedNoSgClusterStorageOk", "attrs.kv_cache_bit_width == 8", "plainAttentionOptions", "(subgroupCachedPrefillCluster and cachedF32ClusterRegisterGeometry) if \"_nosg\" == \"\" else true"], + "constants": { + "scalar": "gqaScalar", + "usesF16": false, + "hasCausal": true, + "causalRightAlign": true, + "hasBias": false, + "hasMask": false, + "maskIsBool": false, + "hasWindow": "attrs.local_window_size > 0", + "headDim": "gqaHeadDim", + "headDimV4": "gqaHeadDim / 4", + "qHiddenV4": "dim(shapes.queryT, 2) / 4", + "qNumHeads": "attrs.num_heads", + "kvNumHeads": "attrs.kv_num_heads", + "TILE_Q": "cachedPrefillQueryTile", + "TILE_K": "cachedNoSgTileK", + "LPQ": "cachedPrefillLanesPerQuery", + "kvHeads": "attrs.kv_num_heads", + "inputVec4": "gqaInputVec4", + "inputScalar": "gqaScalar", + "batchNoSgReduction": "\"_nosg\" == \"_nosg\"", + "useSubgroups": "\"_nosg\" == \"\"", + "quantCacheFormat": "\"int8\"", + "quantCacheElement": "\"i32\"", + "quantCacheReadElement": "\"vec4\"", + "packed": "gqaHeadDim", + "cooperative": false, + "dispatchUnits": 1, + "qHeads": "attrs.num_heads", + "qHidden": "dim(shapes.queryT, 2)", + "mode": "\"build_quant\"", + "bits": 8, + "qmax": 127, + "qmin": -128, + "quantized": true + }, + "passes": [ + { + "id": "present", + "name": "GroupQueryAttention.QuantPresent", + "shader": "gqa-present.wgsl.jinja", + "bindings": "quantPresent", + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)", + "workgroupSize": "copyWorkgroupSize" + } + }, + { + "id": "attention", + "name": "GroupQueryAttention.QuantFlashPrefill", + "source": { + "shader": "attn-flash-prefill-cluster.wgsl.jinja", + "inputs": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" } + }, + "bindings": "flashPrefillQuant", + "dispatch": { + "x": "ceilDiv(dim(shapes.queryT, 1), cachedPrefillQueryTile)", + "y": "attrs.num_heads", + "z": "dim(shapes.queryT, 0)" + } + } + ] + }, + { + "id": "quant_int4_flash_prefill_nosg", + "description": "Builds an int8 or int4 cache and applies clustered online-softmax prefill directly from the packed values. The family supplies subgroup and portable reductions.", + "priority": 37, + "when": ["quantizedPromptOk", "prefillFlashShapeOk", "cachedClusterStorageOk if \"_nosg\" == \"\" else cachedNoSgClusterStorageOk", "attrs.kv_cache_bit_width == 4", "present.kScaleT", "plainAttentionOptions", "(subgroupCachedPrefillCluster and cachedF32ClusterRegisterGeometry) if \"_nosg\" == \"\" else true"], + "constants": { + "scalar": "gqaScalar", + "usesF16": false, + "hasCausal": true, + "causalRightAlign": true, + "hasBias": false, + "hasMask": false, + "maskIsBool": false, + "hasWindow": "attrs.local_window_size > 0", + "headDim": "gqaHeadDim", + "headDimV4": "gqaHeadDim / 4", + "qHiddenV4": "dim(shapes.queryT, 2) / 4", + "qNumHeads": "attrs.num_heads", + "kvNumHeads": "attrs.kv_num_heads", + "TILE_Q": "cachedPrefillQueryTile", + "TILE_K": "cachedNoSgTileK", + "LPQ": "cachedPrefillLanesPerQuery", + "kvHeads": "attrs.kv_num_heads", + "inputVec4": "gqaInputVec4", + "inputScalar": "gqaScalar", + "batchNoSgReduction": "\"_nosg\" == \"_nosg\"", + "useSubgroups": "\"_nosg\" == \"\"", + "quantCacheFormat": "\"int4\"", + "quantCacheElement": "\"u32\"", + "quantCacheReadElement": "\"u32\"", + "packed": "gqaHeadDim / 2", + "cooperative": false, + "dispatchUnits": 1, + "qHeads": "attrs.num_heads", + "qHidden": "dim(shapes.queryT, 2)", + "mode": "\"build_quant\"", + "bits": 4, + "qmax": 7, + "qmin": -8, + "quantized": true + }, + "passes": [ + { + "id": "present", + "name": "GroupQueryAttention.QuantPresent", + "shader": "gqa-present.wgsl.jinja", + "bindings": "quantPresent", + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)", + "workgroupSize": "copyWorkgroupSize" + } + }, + { + "id": "attention", + "name": "GroupQueryAttention.QuantFlashPrefill", + "source": { + "shader": "attn-flash-prefill-cluster.wgsl.jinja", + "inputs": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" } + }, + "bindings": "flashPrefillQuant", + "dispatch": { + "x": "ceilDiv(dim(shapes.queryT, 1), cachedPrefillQueryTile)", + "y": "attrs.num_heads", + "z": "dim(shapes.queryT, 0)" + } + } + ] + }, + { + "id": "share_append_split_decode_splitk", + "description": "Retains shared-capacity cache rows, appends new rows separately, then partitions decode across the key axis. It avoids rebuilding unchanged cache positions while exposing split-key parallelism.", + "priority": 39, + "when": ["shareAppendOk", "decodeFlashShapeOk", "cachedSplitScratchFits", "plainAttentionOptions", "subgroupSplitK if \"\" == \"\" else true"], + "constants": { + "cacheSeqlens": "true", + "scalar": "gqaScalar", + "usesF16": "tensorDtypes.queryT == \"float16\"", + "qNumHeads": "attrs.num_heads", + "kvNumHeads": "attrs.kv_num_heads", + "headDim": "gqaHeadDim", + "headDimV4": "gqaHeadDim / 4", + "numSplits": "cachedNumSplits", + "hasWindow": "attrs.local_window_size > 0", + "kvHeads": "attrs.kv_num_heads", + "packed": "gqaHeadDim", + "mode": "\"merge_share\"", + "inputVec4": "gqaInputVec4", + "inputScalar": "gqaScalar", + "zeroScalar": "gqaScalar", + "combineSubgroups": "\"\" == \"\"", + "useSubgroups": "\"\" == \"\"" + }, + "intermediates": [ + { + "id": "partialOut", + "dtype": "float32", + "shape": "[dim(shapes.queryT, 0) * attrs.num_heads * (cachedNumSplits) * gqaHeadDim]" + }, + { + "id": "partialStats", + "dtype": "float32", + "shape": "[2 * dim(shapes.queryT, 0) * attrs.num_heads * (cachedNumSplits)]" + } + ], + "passes": [ + { + "id": "present_retain", + "name": "GroupQueryAttention.MergeShareRetain", + "shader": "gqa-present.wgsl.jinja", + "bindings": "presentMergeShareRetain", + "constants": { "shareRegion": "\"retain\"" }, + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)", + "workgroupSize": "copyWorkgroupSize" + }, + "viewAlias": [ + { "input": "past_k", "output": "present_key" }, + { "input": "past_v", "output": "present_value" } + ] + }, + { + "id": "present_append", + "name": "GroupQueryAttention.MergeShareAppend", + "shader": "gqa-present.wgsl.jinja", + "bindings": "presentMergeShareAppend", + "constants": { "shareRegion": "\"append\"" }, + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)", + "workgroupSize": "copyWorkgroupSize" + } + }, + { + "id": "split_attention", + "name": "GroupQueryAttention.DecodeSplitK", + "source": { "shader": "attn-flash-decode-splitk.wgsl.jinja", "inputs": { "layout": "\"bhsd\"" } }, + "bindings": "splitSeqlensAttention", + "dispatch": { "x": "cachedNumSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } + }, + { + "id": "merge", + "name": "GroupQueryAttention.DecodeSplitKMerge", + "source": { "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", "inputs": { "layout": "\"bhsd\"" } }, + "bindings": "splitMergeNoParams", + "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } + } + ] + }, + { + "id": "share_append_split_decode_splitk_nosg", + "description": "Retains shared-capacity cache rows, appends new rows separately, then partitions decode across the key axis. It avoids rebuilding unchanged cache positions while exposing split-key parallelism.", + "priority": 39, + "when": ["shareAppendOk", "decodeFlashShapeOk", "cachedSplitScratchFits", "plainAttentionOptions", "subgroupSplitK if \"_nosg\" == \"\" else true"], + "constants": { + "cacheSeqlens": "true", + "scalar": "gqaScalar", + "usesF16": "tensorDtypes.queryT == \"float16\"", + "qNumHeads": "attrs.num_heads", + "kvNumHeads": "attrs.kv_num_heads", + "headDim": "gqaHeadDim", + "headDimV4": "gqaHeadDim / 4", + "numSplits": "cachedNumSplits", + "hasWindow": "attrs.local_window_size > 0", + "kvHeads": "attrs.kv_num_heads", + "packed": "gqaHeadDim", + "mode": "\"merge_share\"", + "inputVec4": "gqaInputVec4", + "inputScalar": "gqaScalar", + "zeroScalar": "gqaScalar", + "combineSubgroups": "\"_nosg\" == \"\"", + "useSubgroups": "\"_nosg\" == \"\"" + }, + "intermediates": [ + { + "id": "partialOut", + "dtype": "float32", + "shape": "[dim(shapes.queryT, 0) * attrs.num_heads * (cachedNumSplits) * gqaHeadDim]" + }, + { + "id": "partialStats", + "dtype": "float32", + "shape": "[2 * dim(shapes.queryT, 0) * attrs.num_heads * (cachedNumSplits)]" + } + ], + "passes": [ + { + "id": "present_retain", + "name": "GroupQueryAttention.MergeShareRetain", + "shader": "gqa-present.wgsl.jinja", + "bindings": "presentMergeShareRetain", + "constants": { "shareRegion": "\"retain\"" }, + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)", + "workgroupSize": "copyWorkgroupSize" + }, + "viewAlias": [ + { "input": "past_k", "output": "present_key" }, + { "input": "past_v", "output": "present_value" } + ] + }, + { + "id": "present_append", + "name": "GroupQueryAttention.MergeShareAppend", + "shader": "gqa-present.wgsl.jinja", + "bindings": "presentMergeShareAppend", + "constants": { "shareRegion": "\"append\"" }, + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)", + "workgroupSize": "copyWorkgroupSize" + } + }, + { + "id": "split_attention", + "name": "GroupQueryAttention.DecodeSplitK", + "source": { "shader": "attn-flash-decode-splitk.wgsl.jinja", "inputs": { "layout": "\"bhsd\"" } }, + "bindings": "splitSeqlensAttention", + "dispatch": { "x": "cachedNumSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } + }, + { + "id": "merge", + "name": "GroupQueryAttention.DecodeSplitKMerge", + "source": { "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", "inputs": { "layout": "\"bhsd\"" } }, + "bindings": "splitMergeNoParams", + "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } + } + ] + }, + { + "id": "share_append_split_flash_prefill", + "description": "Retains shared-capacity cache rows, appends new rows separately, then applies clustered causal prefill. It avoids rebuilding unchanged cache positions and includes subgroup and portable forms.", + "priority": 35, + "when": ["shareAppendOk", "prefillFlashShapeOk", "cachedClusterStorageOk if \"\" == \"\" else cachedNoSgClusterStorageOk", "plainAttentionOptions", "(subgroupCachedPrefillCluster and cachedF32ClusterRegisterGeometry) if \"\" == \"\" else true"], + "constants": { + "scalar": "gqaScalar", + "usesF16": "tensorDtypes.queryT == \"float16\"", + "hasCausal": true, + "causalRightAlign": true, + "hasBias": false, + "hasMask": false, + "maskIsBool": false, + "hasWindow": "attrs.local_window_size > 0", + "headDim": "gqaHeadDim", + "headDimV4": "gqaHeadDim / 4", + "qHiddenV4": "dim(shapes.queryT, 2) / 4", + "qNumHeads": "attrs.num_heads", + "kvNumHeads": "attrs.kv_num_heads", + "TILE_Q": "cachedPrefillQueryTile", + "TILE_K": "cachedClusterTileK", + "LPQ": "cachedPrefillLanesPerQuery", + "kvHeads": "attrs.kv_num_heads", + "packed": "gqaHeadDim", + "mode": "\"merge_share\"", + "useSeqlens": true, + "inputVec4": "gqaInputVec4", + "inputScalar": "gqaScalar", + "batchNoSgReduction": "\"\" == \"_nosg\"", + "useSubgroups": "\"\" == \"\"" + }, + "passes": [ + { + "id": "present_retain", + "name": "GroupQueryAttention.MergeShareRetain", + "shader": "gqa-present.wgsl.jinja", + "bindings": "presentMergeShareRetain", + "constants": { "shareRegion": "\"retain\"" }, + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)", + "workgroupSize": "copyWorkgroupSize" + }, + "viewAlias": [ + { "input": "past_k", "output": "present_key" }, + { "input": "past_v", "output": "present_value" } + ] + }, + { + "id": "present_append", + "name": "GroupQueryAttention.MergeShareAppend", + "shader": "gqa-present.wgsl.jinja", + "bindings": "presentMergeShareAppend", + "constants": { "shareRegion": "\"append\"" }, + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)", + "workgroupSize": "copyWorkgroupSize" + } + }, + { + "id": "attention", + "name": "GroupQueryAttention.FlashPrefill", + "source": { + "shader": "attn-flash-prefill-cluster.wgsl.jinja", + "inputs": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" } + }, + "bindings": "flashPrefillSeqlens", + "dispatch": { + "x": "ceilDiv(dim(shapes.queryT, 1), constants.TILE_Q)", + "y": "attrs.num_heads", + "z": "dim(shapes.queryT, 0)" + } + } + ] + }, + { + "id": "share_append_split_flash_prefill_nosg", + "description": "Retains shared-capacity cache rows, appends new rows separately, then applies clustered causal prefill. It avoids rebuilding unchanged cache positions and includes subgroup and portable forms.", + "priority": 35, + "when": ["shareAppendOk", "prefillFlashShapeOk", "cachedClusterStorageOk if \"_nosg\" == \"\" else cachedNoSgClusterStorageOk", "plainAttentionOptions", "(subgroupCachedPrefillCluster and cachedF32ClusterRegisterGeometry) if \"_nosg\" == \"\" else true"], + "constants": { + "scalar": "gqaScalar", + "usesF16": "tensorDtypes.queryT == \"float16\"", + "hasCausal": true, + "causalRightAlign": true, + "hasBias": false, + "hasMask": false, + "maskIsBool": false, + "hasWindow": "attrs.local_window_size > 0", + "headDim": "gqaHeadDim", + "headDimV4": "gqaHeadDim / 4", + "qHiddenV4": "dim(shapes.queryT, 2) / 4", + "qNumHeads": "attrs.num_heads", + "kvNumHeads": "attrs.kv_num_heads", + "TILE_Q": "cachedPrefillQueryTile", + "TILE_K": "cachedNoSgTileK", + "LPQ": "cachedPrefillLanesPerQuery", + "kvHeads": "attrs.kv_num_heads", + "packed": "gqaHeadDim", + "mode": "\"merge_share\"", + "useSeqlens": true, + "inputVec4": "gqaInputVec4", + "inputScalar": "gqaScalar", + "batchNoSgReduction": "\"_nosg\" == \"_nosg\"", + "useSubgroups": "\"_nosg\" == \"\"" + }, + "passes": [ + { + "id": "present_retain", + "name": "GroupQueryAttention.MergeShareRetain", + "shader": "gqa-present.wgsl.jinja", + "bindings": "presentMergeShareRetain", + "constants": { "shareRegion": "\"retain\"" }, + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)", + "workgroupSize": "copyWorkgroupSize" + }, + "viewAlias": [ + { "input": "past_k", "output": "present_key" }, + { "input": "past_v", "output": "present_value" } + ] + }, + { + "id": "present_append", + "name": "GroupQueryAttention.MergeShareAppend", + "shader": "gqa-present.wgsl.jinja", + "bindings": "presentMergeShareAppend", + "constants": { "shareRegion": "\"append\"" }, + "dispatch": { + "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)", + "workgroupSize": "copyWorkgroupSize" + } + }, + { + "id": "attention", + "name": "GroupQueryAttention.FlashPrefill", + "source": { + "shader": "attn-flash-prefill-cluster.wgsl.jinja", + "inputs": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" } + }, + "bindings": "flashPrefillSeqlens", + "dispatch": { + "x": "ceilDiv(dim(shapes.queryT, 1), constants.TILE_Q)", + "y": "attrs.num_heads", + "z": "dim(shapes.queryT, 0)" + } + } + ] + } + ] +}