| { |
| "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, |
| "MATERIALIZED_SGMAT_QUERY_TILE": 64, |
| "MATERIALIZED_SGMAT_KEY_TILE": 64, |
| "MATERIALIZED_SGMAT_INNER_TILE": 32, |
| "MATERIALIZED_SOFTMAX_WORKGROUP_SIZE": 256, |
| "MATERIALIZED_SGMAT_MIN_SEQ": 256, |
| "MATERIALIZED_SGMAT_PAST_MIN_SEQ": 256, |
| "MATERIALIZED_SGMAT_WINDOW_MIN_QSEQ": 128 |
| }, |
| "derive": { |
| "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)", |
| "wave32Adapter": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize == 32 and device.adapterInfo.subgroupMaxSize == 32", |
| "subgroupsWave32": "device.features.has(\"subgroups\") and wave32Adapter", |
| "narrowSubgroupRange": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize < device.adapterInfo.subgroupMaxSize and device.adapterInfo.subgroupMaxSize <= 16", |
| "canPinSubgroupSize32": "device.features.has(\"subgroups\") and device.features.has(\"subgroup-size-control\") and has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize <= 32 and device.adapterInfo.subgroupMaxSize >= 32", |
| "pinSubgroupSize32": "canPinSubgroupSize32 and not wave32Adapter", |
| "wave32Effective": "wave32Adapter or pinSubgroupSize32", |
| "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<f16>\" if tensorDtypes.query == \"float16\" else \"vec4<f32>\"", |
| "qkvCacheScalar": "\"f16\" if tensorDtypes.presentKeyT == \"float16\" else \"f32\"", |
| "qkvCacheVec4": "\"vec4<f16>\" if tensorDtypes.presentKeyT == \"float16\" else \"vec4<f32>\"", |
| "qkvUsesF16": "tensorDtypes.query == \"float16\" or tensorDtypes.presentKeyT == \"float16\"", |
| "gqaScalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"", |
| "gqaInputVec4": "\"vec4<f16>\" if tensorDtypes.queryT == \"float16\" else \"vec4<f32>\"", |
| "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", |
| "gqaMatQueryTile": "tunables.MATERIALIZED_SGMAT_QUERY_TILE", |
| "gqaMatKeyTile": "tunables.MATERIALIZED_SGMAT_KEY_TILE", |
| "gqaMatInnerTile": "tunables.MATERIALIZED_SGMAT_INNER_TILE", |
| "gqaMatSubgroupRows": "floor(gqaMatQueryTile / 16)", |
| "gqaMatSubgroupCols": "floor(gqaMatKeyTile / 32)", |
| "gqaMatWorkgroupSize": "gqaMatSubgroupRows * gqaMatSubgroupCols * 32", |
| "gqaMatCompactStorageBytes": "(gqaMatQueryTile + gqaMatKeyTile) * gqaMatInnerTile * 4", |
| "gqaMatGeometryOk": "gqaMatQueryTile >= 16 and gqaMatQueryTile % 16 == 0 and gqaMatKeyTile >= 32 and gqaMatKeyTile <= 64 and gqaMatKeyTile % 32 == 0 and gqaMatInnerTile == 32 and gqaMatQueryTile >= 64 and gqaMatKeyTile >= 64", |
| "gqaMatResourcesFit": "gqaMatGeometryOk and gqaMatWorkgroupSize <= deviceWorkgroupCap and gqaMatCompactStorageBytes <= device.limits.maxComputeWorkgroupStorageSize", |
| "gqaMatScoreBytes": "dim(shapes.query, 0) * attrs.num_heads * dim(shapes.query, 1) * dim(shapes.key, 1) * 4", |
| "gqaMatScoreFits": "gqaMatScoreBytes <= device.limits.maxStorageBufferBindingSize and gqaMatScoreBytes <= device.limits.maxBufferSize", |
| "gqaMatBuffersFit": "numel(shapes.query) * 4 <= device.limits.maxStorageBufferBindingSize and numel(shapes.key) * 4 <= device.limits.maxStorageBufferBindingSize and numel(shapes.value) * 4 <= device.limits.maxStorageBufferBindingSize", |
| "gqaMatDispatchFits": "ceilDiv(dim(shapes.key, 1), gqaMatKeyTile) <= device.limits.maxComputeWorkgroupsPerDimension and ceilDiv(dim(shapes.query, 1), gqaMatQueryTile) <= device.limits.maxComputeWorkgroupsPerDimension and dim(shapes.query, 0) * attrs.num_heads <= device.limits.maxComputeWorkgroupsPerDimension", |
| "gqaMatStatSlots": "ceilDiv(dim(shapes.key, 1), gqaMatKeyTile) * gqaMatSubgroupCols", |
| "gqaMatRowStatsWg": "min(tunables.MATERIALIZED_SOFTMAX_WORKGROUP_SIZE, deviceWorkgroupCap)", |
| "gqaMatRowStatsElements": "dim(shapes.query, 0) * attrs.num_heads * dim(shapes.query, 1) * 2", |
| "gqaMatScorePartialElements": "dim(shapes.query, 0) * attrs.num_heads * gqaMatStatSlots * dim(shapes.query, 1) * 2", |
| "gqaMatCoreOk": "qkvContractOk and attrs.causal == 0 and tensorDtypes.query == \"float32\" and attrs.local_window_size == -1 and headDim >= 64 and headDim <= 256 and headDim % 16 == 0 and dim(shapes.query, 1) >= tunables.MATERIALIZED_SGMAT_MIN_SEQ and dim(shapes.key, 1) >= tunables.MATERIALIZED_SGMAT_MIN_SEQ and device.features.has(\"subgroups\") and wave32Effective and device.features.has(\"chromium-experimental-subgroup-matrix\") and gqaMatScoreFits and gqaMatBuffersFit and gqaMatResourcesFit and gqaMatDispatchFits", |
| "gqaMatPastQSeq": "dim(shapes.queryT, 1)", |
| "gqaMatPastKvSeq": "dim(shapes.presentKeyT, 2)", |
| "gqaMatPastScoreBytes": "dim(shapes.queryT, 0) * attrs.num_heads * gqaMatPastQSeq * gqaMatPastKvSeq * 4", |
| "gqaMatPastScoreFits": "gqaMatPastScoreBytes <= device.limits.maxStorageBufferBindingSize and gqaMatPastScoreBytes <= device.limits.maxBufferSize", |
| "gqaMatPastBuffersFit": "numel(shapes.queryT) * 4 <= device.limits.maxStorageBufferBindingSize and numel(shapes.presentKeyT) * 4 <= device.limits.maxStorageBufferBindingSize and numel(shapes.presentValueT) * 4 <= device.limits.maxStorageBufferBindingSize", |
| "gqaMatPastStatSlots": "ceilDiv(gqaMatPastKvSeq, gqaMatKeyTile) * gqaMatSubgroupCols", |
| "gqaMatPastRowStatsElements": "dim(shapes.queryT, 0) * attrs.num_heads * gqaMatPastQSeq * 2", |
| "gqaMatPastScorePartialElements": "dim(shapes.queryT, 0) * attrs.num_heads * gqaMatPastStatSlots * gqaMatPastQSeq * 2", |
| "gqaMatPastDispatchFits": "ceilDiv(gqaMatPastKvSeq, gqaMatKeyTile) <= device.limits.maxComputeWorkgroupsPerDimension and ceilDiv(gqaMatPastQSeq, gqaMatQueryTile) <= device.limits.maxComputeWorkgroupsPerDimension and dim(shapes.queryT, 0) * attrs.num_heads <= device.limits.maxComputeWorkgroupsPerDimension", |
| "gqaMatWindowOk": "windowShiftOk and plainAttentionOptions and tensorDtypes.queryT == \"float32\" and attrs.local_window_size > 0 and dim(shapes.queryT, 1) == dim(shapes.keyT, 1) and gqaHeadDim >= 64 and gqaHeadDim <= 256 and gqaHeadDim % 16 == 0 and dim(shapes.queryT, 1) >= tunables.MATERIALIZED_SGMAT_WINDOW_MIN_QSEQ and windowCapacity >= dim(shapes.queryT, 1) and attrs.local_window_size + dim(shapes.queryT, 1) >= windowCapacity and device.features.has(\"subgroups\") and wave32Effective and device.features.has(\"chromium-experimental-subgroup-matrix\") and gqaMatPastScoreFits and gqaMatPastBuffersFit and gqaMatResourcesFit and gqaMatPastDispatchFits", |
| "gqaMatPastCoreOk": "sharedKvFloatOk and plainAttentionOptions and tensorDtypes.queryT == \"float32\" and attrs.local_window_size == -1 and gqaHeadDim >= 64 and gqaHeadDim <= 256 and gqaHeadDim % 16 == 0 and gqaMatPastQSeq >= tunables.MATERIALIZED_SGMAT_PAST_MIN_SEQ and gqaMatPastKvSeq >= tunables.MATERIALIZED_SGMAT_PAST_MIN_SEQ and gqaMatPastKvSeq >= gqaMatPastQSeq and device.features.has(\"subgroups\") and wave32Effective and device.features.has(\"chromium-experimental-subgroup-matrix\") and gqaMatPastScoreFits and gqaMatPastBuffersFit and gqaMatResourcesFit and gqaMatPastDispatchFits", |
| "gqaMatPastRotaryOk": "sharedKvFloatOk and standardSoftmax and present.cosCacheT and rotaryRequested and not present.qNormWeightT and not present.attentionBiasT and not present.headSinkT and tensorDtypes.queryT == \"float32\" and attrs.local_window_size == -1 and gqaHeadDim >= 64 and gqaHeadDim <= 256 and gqaHeadDim % 16 == 0 and gqaHeadDim % 8 == 0 and gqaMatPastQSeq >= tunables.MATERIALIZED_SGMAT_PAST_MIN_SEQ and gqaMatPastKvSeq >= tunables.MATERIALIZED_SGMAT_PAST_MIN_SEQ and gqaMatPastKvSeq >= gqaMatPastQSeq and device.features.has(\"subgroups\") and wave32Effective and device.features.has(\"chromium-experimental-subgroup-matrix\") and gqaMatPastScoreFits and gqaMatPastBuffersFit and gqaMatResourcesFit and gqaMatPastDispatchFits", |
| "gqaMatNewPastOk": "newKvPastOk and plainAttentionOptions and tensorDtypes.queryT == \"float32\" and attrs.local_window_size == -1 and gqaHeadDim >= 64 and gqaHeadDim <= 256 and gqaHeadDim % 16 == 0 and gqaMatPastQSeq >= tunables.MATERIALIZED_SGMAT_PAST_MIN_SEQ and gqaMatPastKvSeq >= tunables.MATERIALIZED_SGMAT_PAST_MIN_SEQ and gqaMatPastKvSeq >= gqaMatPastQSeq and device.features.has(\"subgroups\") and wave32Effective and device.features.has(\"chromium-experimental-subgroup-matrix\") and gqaMatPastScoreFits and gqaMatPastBuffersFit and gqaMatResourcesFit and gqaMatPastDispatchFits", |
| "gqaMatPastCoreF16Ok": "sharedKvFloatOk and plainAttentionOptions and tensorDtypes.queryT == \"float16\" and device.features.has(\"shader-f16\") and attrs.local_window_size == -1 and gqaHeadDim >= 64 and gqaHeadDim <= 256 and gqaHeadDim % 16 == 0 and gqaMatPastQSeq >= tunables.MATERIALIZED_SGMAT_PAST_MIN_SEQ and gqaMatPastKvSeq >= tunables.MATERIALIZED_SGMAT_PAST_MIN_SEQ and gqaMatPastKvSeq >= gqaMatPastQSeq and device.features.has(\"subgroups\") and wave32Effective and device.features.has(\"chromium-experimental-subgroup-matrix\") and gqaMatPastScoreFits and gqaMatPastBuffersFit and gqaMatResourcesFit and gqaMatPastDispatchFits", |
| "gqaMatPastRotaryF16Ok": "sharedKvFloatOk and standardSoftmax and present.cosCacheT and rotaryRequested and not present.qNormWeightT and not present.attentionBiasT and not present.headSinkT and tensorDtypes.queryT == \"float16\" and device.features.has(\"shader-f16\") and attrs.local_window_size == -1 and gqaHeadDim >= 64 and gqaHeadDim <= 256 and gqaHeadDim % 16 == 0 and gqaHeadDim % 8 == 0 and gqaMatPastQSeq >= tunables.MATERIALIZED_SGMAT_PAST_MIN_SEQ and gqaMatPastKvSeq >= tunables.MATERIALIZED_SGMAT_PAST_MIN_SEQ and gqaMatPastKvSeq >= gqaMatPastQSeq and device.features.has(\"subgroups\") and wave32Effective and device.features.has(\"chromium-experimental-subgroup-matrix\") and gqaMatPastScoreFits and gqaMatPastBuffersFit and gqaMatResourcesFit and gqaMatPastDispatchFits", |
| "gqaMatPastBiasOk": "sharedKvFloatOk and standardSoftmax and not present.cosCacheT and not rotaryRequested and not present.qNormWeightT and not present.headSinkT and present.attentionBiasT and ranks.attentionBiasT == 4 and tensorDtypes.attentionBiasT == \"float32\" and (dim(shapes.attentionBiasT, 0) == dim(shapes.queryT, 0) or dim(shapes.attentionBiasT, 0) == 1) and (dim(shapes.attentionBiasT, 1) == attrs.num_heads or dim(shapes.attentionBiasT, 1) == 1) and dim(shapes.attentionBiasT, 2) == dim(shapes.queryT, 1) and dim(shapes.attentionBiasT, 3) == dim(shapes.presentKeyT, 2) and tensorDtypes.queryT == \"float32\" and attrs.local_window_size == -1 and gqaHeadDim >= 64 and gqaHeadDim <= 256 and gqaHeadDim % 16 == 0 and gqaMatPastQSeq >= tunables.MATERIALIZED_SGMAT_PAST_MIN_SEQ and gqaMatPastKvSeq >= tunables.MATERIALIZED_SGMAT_PAST_MIN_SEQ and gqaMatPastKvSeq >= gqaMatPastQSeq and device.features.has(\"subgroups\") and wave32Effective and device.features.has(\"chromium-experimental-subgroup-matrix\") and gqaMatPastScoreFits and gqaMatPastBuffersFit and gqaMatResourcesFit and gqaMatPastDispatchFits" |
| }, |
| "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<f32>" }, |
| { |
| "name": "partial_stats", |
| "semantic": "partialStats", |
| "buffer": { "type": "storage" }, |
| "elementType": "vec2<f32>" |
| }, |
| { |
| "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<f32>" }, |
| { |
| "name": "partial_stats", |
| "semantic": "partialStats", |
| "buffer": { "type": "storage" }, |
| "elementType": "vec2<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": "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<f32>" |
| }, |
| { |
| "name": "partial_stats", |
| "semantic": "partialStats", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "vec2<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" }] |
| } |
| } |
| ], |
| "splitMergeNoParams": [ |
| { |
| "name": "partial_out", |
| "semantic": "partialOut", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "vec4<f32>" |
| }, |
| { |
| "name": "partial_stats", |
| "semantic": "partialStats", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "vec2<f32>" |
| }, |
| { "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<f32>" }, |
| { "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<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.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<f32>" }, |
| { |
| "name": "partial_stats", |
| "semantic": "partialStats", |
| "buffer": { "type": "storage" }, |
| "elementType": "vec2<f32>" |
| }, |
| { |
| "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<f32>" |
| }, |
| { |
| "name": "partial_stats", |
| "semantic": "partialStats", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "vec2<f32>" |
| }, |
| { |
| "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<f32>" |
| }, |
| { |
| "name": "partial_stats", |
| "semantic": "partialStats", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "vec2<f32>" |
| }, |
| { |
| "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<f32>" }, |
| { "name": "key", "arg": "presentKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "vec4<i32>" }, |
| { "name": "value", "arg": "presentValueT", "buffer": { "type": "read-only-storage" }, "elementType": "vec4<i32>" }, |
| { "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<f32>" }, |
| { |
| "name": "partial_stats", |
| "semantic": "partialStats", |
| "buffer": { "type": "storage" }, |
| "elementType": "vec2<f32>" |
| }, |
| { |
| "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" } |
| ] |
| } |
| } |
| ], |
| "gqaMatScoreStats": [ |
| { |
| "name": "query", |
| "arg": "queryT", |
| "semantic": "query", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "key", |
| "arg": "keyT", |
| "semantic": "key", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { "name": "scores", "semantic": "materializedScores", "buffer": { "type": "storage" }, "elementType": "f32" }, |
| { |
| "name": "scorePartials", |
| "semantic": "materializedScorePartials", |
| "buffer": { "type": "storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "qSeq", "type": "u32", "value": "dim(shapes.query, 1)" }, |
| { "name": "kvSeq", "type": "u32", "value": "dim(shapes.key, 1)" }, |
| { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" } |
| ] |
| } |
| } |
| ], |
| "gqaMatRowStatsCombine": [ |
| { |
| "name": "scorePartials", |
| "semantic": "materializedScorePartials", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { "name": "rowStats", "semantic": "materializedRowStats", "buffer": { "type": "storage" }, "elementType": "f32" }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "rows", "type": "u32", "value": "dim(shapes.query, 0) * attrs.num_heads * dim(shapes.query, 1)" } |
| ] |
| } |
| } |
| ], |
| "gqaMatApplyFused": [ |
| { |
| "name": "scores", |
| "semantic": "materializedScores", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "value", |
| "arg": "valueT", |
| "semantic": "value", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "rowStats", |
| "semantic": "materializedRowStats", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "output", |
| "arg": "outputT", |
| "semantic": "output", |
| "buffer": { "type": "storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "qSeq", "type": "u32", "value": "dim(shapes.query, 1)" }, |
| { "name": "kvSeq", "type": "u32", "value": "dim(shapes.key, 1)" } |
| ] |
| } |
| } |
| ], |
| "gqaMatPastScoreStats": [ |
| { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, |
| { "name": "key", "arg": "presentKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, |
| { "name": "scores", "semantic": "materializedScores", "buffer": { "type": "storage" }, "elementType": "f32" }, |
| { |
| "name": "scorePartials", |
| "semantic": "materializedScorePartials", |
| "buffer": { "type": "storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, |
| { "name": "kvSeq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" }, |
| { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" } |
| ] |
| } |
| } |
| ], |
| "gqaMatWindowScoreStats": [ |
| { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, |
| { "name": "key", "arg": "presentKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, |
| { "name": "scores", "semantic": "materializedScores", "buffer": { "type": "storage" }, "elementType": "f32" }, |
| { |
| "name": "scorePartials", |
| "semantic": "materializedScorePartials", |
| "buffer": { "type": "storage" }, |
| "elementType": "f32" |
| }, |
| { "name": "seqlens_k", "arg": "seqlensKT", "buffer": { "type": "read-only-storage" }, "elementType": "i32" }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, |
| { "name": "kvSeq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" }, |
| { "name": "windowSize", "type": "u32", "value": "attrs.local_window_size" }, |
| { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" } |
| ] |
| } |
| } |
| ], |
| "gqaMatWindowApplyFused": [ |
| { |
| "name": "scores", |
| "semantic": "materializedScores", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { "name": "value", "arg": "presentValueT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, |
| { |
| "name": "rowStats", |
| "semantic": "materializedRowStats", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { "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": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, |
| { "name": "kvSeq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" }, |
| { "name": "windowSize", "type": "u32", "value": "attrs.local_window_size" } |
| ] |
| } |
| } |
| ], |
| "gqaMatPastRotaryScoreStats": [ |
| { "name": "query", "semantic": "qPrep", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, |
| { "name": "key", "arg": "presentKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, |
| { "name": "scores", "semantic": "materializedScores", "buffer": { "type": "storage" }, "elementType": "f32" }, |
| { |
| "name": "scorePartials", |
| "semantic": "materializedScorePartials", |
| "buffer": { "type": "storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, |
| { "name": "kvSeq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" }, |
| { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" } |
| ] |
| } |
| } |
| ], |
| "rotaryQprep": [ |
| { "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": "qout", "semantic": "qPrep", "buffer": { "type": "storage" }, "elementType": "vec4<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.pastKeyT, 2)" } |
| ] |
| } |
| } |
| ], |
| "gqaMatPastRowStatsCombine": [ |
| { |
| "name": "scorePartials", |
| "semantic": "materializedScorePartials", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { "name": "rowStats", "semantic": "materializedRowStats", "buffer": { "type": "storage" }, "elementType": "f32" }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { |
| "name": "rows", |
| "type": "u32", |
| "value": "dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)" |
| } |
| ] |
| } |
| } |
| ], |
| "gqaMatPastApplyFused": [ |
| { |
| "name": "scores", |
| "semantic": "materializedScores", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { "name": "value", "arg": "presentValueT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, |
| { |
| "name": "rowStats", |
| "semantic": "materializedRowStats", |
| "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": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, |
| { "name": "kvSeq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" } |
| ] |
| } |
| } |
| ], |
| "gqaMatPastScoreStatsF16": [ |
| { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "f16" }, |
| { "name": "key", "arg": "presentKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "f16" }, |
| { "name": "scores", "semantic": "materializedScores", "buffer": { "type": "storage" }, "elementType": "f32" }, |
| { |
| "name": "scorePartials", |
| "semantic": "materializedScorePartials", |
| "buffer": { "type": "storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, |
| { "name": "kvSeq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" }, |
| { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" } |
| ] |
| } |
| } |
| ], |
| "gqaMatPastApplyFusedF16": [ |
| { |
| "name": "scores", |
| "semantic": "materializedScores", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { "name": "value", "arg": "presentValueT", "buffer": { "type": "read-only-storage" }, "elementType": "f16" }, |
| { |
| "name": "rowStats", |
| "semantic": "materializedRowStats", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { "name": "output", "arg": "outputT", "buffer": { "type": "storage" }, "elementType": "f16" }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, |
| { "name": "kvSeq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" } |
| ] |
| } |
| } |
| ], |
| "gqaMatPastRotaryScoreStatsF16": [ |
| { "name": "query", "semantic": "qPrep", "buffer": { "type": "read-only-storage" }, "elementType": "f16" }, |
| { "name": "key", "arg": "presentKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "f16" }, |
| { "name": "scores", "semantic": "materializedScores", "buffer": { "type": "storage" }, "elementType": "f32" }, |
| { |
| "name": "scorePartials", |
| "semantic": "materializedScorePartials", |
| "buffer": { "type": "storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, |
| { "name": "kvSeq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" }, |
| { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" } |
| ] |
| } |
| } |
| ], |
| "rotaryQprepF16": [ |
| { "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": "qout", "semantic": "qPrep", "buffer": { "type": "storage" }, "elementType": "vec4<f16>" }, |
| { |
| "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)" } |
| ] |
| } |
| } |
| ], |
| "gqaMatPastBiasScoreStats": [ |
| { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, |
| { "name": "key", "arg": "presentKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, |
| { |
| "name": "attention_bias", |
| "arg": "attentionBiasT", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { "name": "scores", "semantic": "materializedScores", "buffer": { "type": "storage" }, "elementType": "f32" }, |
| { |
| "name": "scorePartials", |
| "semantic": "materializedScorePartials", |
| "buffer": { "type": "storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, |
| { "name": "kvSeq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" }, |
| { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, |
| { "name": "biasBatch", "type": "u32", "value": "dim(shapes.attentionBiasT, 0)" }, |
| { "name": "biasHeads", "type": "u32", "value": "dim(shapes.attentionBiasT, 1)" } |
| ] |
| } |
| } |
| ] |
| }, |
| "variants": [ |
| { |
| "id": "qkv_present_materialized_sgmat_f32", |
| "description": "Materialized float32 subgroup-matrix prefill for the bidirectional no-past qkv route, where `seqlens_k` is metadata-only like the flash routes: the score pass emits per-row softmax statistics, the apply pass folds the softmax, and the present copy is unchanged. It serves square-ish f32 prefill where the score and value GEMMs dominate.", |
| "priority": 23, |
| "requires": { |
| "features": ["subgroups", "chromium-experimental-subgroup-matrix"], |
| "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] |
| }, |
| "when": ["gqaMatCoreOk"], |
| "constants": { |
| "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)", |
| "hasBias": false, |
| "useSubgroups": true, |
| "attentionScaleExpression": "\"select(inverseSqrt(f32(HEAD_DIM)), params.scale, params.scale != 0.0)\"", |
| "materializedSgmatQueryTile": "gqaMatQueryTile", |
| "materializedSgmatKeyTile": "gqaMatKeyTile", |
| "materializedSgmatInnerTile": "gqaMatInnerTile", |
| "materializedSgmatDirectScoreStore": false, |
| "materializedSgmatDirectApplyStore": false, |
| "materializedSgmatRuntimeDirectStore": false, |
| "materializedRowStatsWg": "gqaMatRowStatsWg", |
| "statSlots": "gqaMatStatSlots", |
| "statQuerySeq": "dim(shapes.query, 1)", |
| "presentScalar": "qkvCacheScalar", |
| "presentElement": "qkvCacheVec4", |
| "presentVec4": true, |
| "copyWorkgroupSize": "copyWorkgroupSize", |
| "inputElement": "qkvInputVec4", |
| "headDimV4": "(dim(shapes.query, 2) / attrs.num_heads) / 4", |
| "kvHiddenV4": "dim(shapes.key, 2) / 4" |
| }, |
| "intermediates": [ |
| { |
| "id": "materializedScores", |
| "dtype": "float32", |
| "shape": "[dim(shapes.query, 0) * attrs.num_heads * dim(shapes.query, 1) * dim(shapes.key, 1)]" |
| }, |
| { "id": "materializedRowStats", "dtype": "float32", "shape": "[gqaMatRowStatsElements]" }, |
| { "id": "materializedScorePartials", "dtype": "float32", "shape": "[gqaMatScorePartialElements]" } |
| ], |
| "passes": [ |
| { |
| "id": "scores", |
| "name": "GroupQueryAttention.MaterializedScoresSgmat", |
| "source": { |
| "shader": "attn-materialized-sgmat-f32.wgsl.jinja", |
| "inputs": { "phase": "\"score\"", "emitRowStats": true, "layout": "\"bsh\"" } |
| }, |
| "bindings": "gqaMatScoreStats", |
| "dispatch": { |
| "x": "ceilDiv(dim(shapes.key, 1), gqaMatKeyTile)", |
| "y": "ceilDiv(dim(shapes.query, 1), gqaMatQueryTile)", |
| "z": "dim(shapes.query, 0) * attrs.num_heads" |
| } |
| }, |
| { |
| "id": "rowstats", |
| "name": "GroupQueryAttention.MaterializedRowStatsCombine", |
| "source": { "shader": "attn-materialized-rowstats-combine-f32.wgsl.jinja" }, |
| "bindings": "gqaMatRowStatsCombine", |
| "dispatch": { |
| "gridStride": "dim(shapes.query, 0) * attrs.num_heads * dim(shapes.query, 1)", |
| "workgroupSize": "gqaMatRowStatsWg" |
| } |
| }, |
| { |
| "id": "apply", |
| "name": "GroupQueryAttention.MaterializedApplySgmat", |
| "source": { |
| "shader": "attn-materialized-sgmat-f32.wgsl.jinja", |
| "inputs": { "phase": "\"apply\"", "fusedSoftmax": true, "layout": "\"bsh\"" } |
| }, |
| "bindings": "gqaMatApplyFused", |
| "dispatch": { |
| "x": "ceilDiv(headDim, gqaMatKeyTile)", |
| "y": "ceilDiv(dim(shapes.query, 1), gqaMatQueryTile)", |
| "z": "dim(shapes.query, 0) * attrs.num_heads" |
| } |
| }, |
| { |
| "id": "present", |
| "name": "GroupQueryAttention.MaterializedPresent", |
| "shader": "gqa-present.wgsl.jinja", |
| "bindings": "qkvPresent", |
| "dispatch": { "threads": "numel(shapes.presentKeyT) / 4", "workgroupSize": "constants.copyWorkgroupSize" } |
| } |
| ] |
| }, |
| { |
| "id": "past_kv_materialized_sgmat_f32", |
| "requires": { |
| "features": ["subgroups", "chromium-experimental-subgroup-matrix"], |
| "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] |
| }, |
| "description": "Materializes the causal score matrix with float32 subgroup-matrix tiles over the shared float cache and applies the softmax-normalized weights with the same tiles. The score pass skips key tiles past each query tile's causal bound and the apply pass stops its reduction there, matching flash's triangular work at matrix-unit throughput.", |
| "priority": 35, |
| "when": ["gqaMatPastCoreOk"], |
| "constants": { |
| "qNumHeads": "attrs.num_heads", |
| "kvNumHeads": "attrs.kv_num_heads", |
| "headDim": "gqaHeadDim", |
| "qHidden": "dim(shapes.queryT, 2)", |
| "hasBias": false, |
| "useSubgroups": true, |
| "attentionScaleExpression": "\"select(inverseSqrt(f32(HEAD_DIM)), params.scale, params.scale != 0.0)\"", |
| "materializedSgmatQueryTile": "gqaMatQueryTile", |
| "materializedSgmatKeyTile": "gqaMatKeyTile", |
| "materializedSgmatInnerTile": "gqaMatInnerTile", |
| "materializedSgmatDirectScoreStore": false, |
| "materializedSgmatDirectApplyStore": false, |
| "materializedSgmatRuntimeDirectStore": false, |
| "materializedRowStatsWg": "gqaMatRowStatsWg", |
| "statSlots": "gqaMatPastStatSlots", |
| "statQuerySeq": "dim(shapes.queryT, 1)", |
| "mode": "\"copy\"", |
| "packed": "gqaHeadDim", |
| "kvHeads": "attrs.kv_num_heads", |
| "inputScalar": "gqaScalar", |
| "usesF16": false, |
| "copyWorkgroupSize": "copyWorkgroupSize" |
| }, |
| "intermediates": [ |
| { |
| "id": "materializedScores", |
| "dtype": "float32", |
| "shape": "[dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1) * dim(shapes.presentKeyT, 2)]" |
| }, |
| { "id": "materializedRowStats", "dtype": "float32", "shape": "[gqaMatPastRowStatsElements]" }, |
| { "id": "materializedScorePartials", "dtype": "float32", "shape": "[gqaMatPastScorePartialElements]" } |
| ], |
| "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": "scores", |
| "name": "GroupQueryAttention.PastMaterializedScores", |
| "source": { |
| "shader": "attn-materialized-sgmat-f32.wgsl.jinja", |
| "inputs": { |
| "phase": "\"score\"", |
| "emitRowStats": true, |
| "layout": "\"bsh\"", |
| "kvLayout": "\"bhsd\"", |
| "causalRightAlign": true |
| } |
| }, |
| "bindings": "gqaMatPastScoreStats", |
| "dispatch": { |
| "x": "ceilDiv(dim(shapes.presentKeyT, 2), gqaMatKeyTile)", |
| "y": "ceilDiv(dim(shapes.queryT, 1), gqaMatQueryTile)", |
| "z": "dim(shapes.queryT, 0) * attrs.num_heads" |
| } |
| }, |
| { |
| "id": "rowstats", |
| "name": "GroupQueryAttention.PastMaterializedRowStats", |
| "source": { "shader": "attn-materialized-rowstats-combine-f32.wgsl.jinja", "inputs": {} }, |
| "bindings": "gqaMatPastRowStatsCombine", |
| "dispatch": { |
| "gridStride": "dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)", |
| "workgroupSize": "gqaMatRowStatsWg" |
| } |
| }, |
| { |
| "id": "apply", |
| "name": "GroupQueryAttention.PastMaterializedApply", |
| "source": { |
| "shader": "attn-materialized-sgmat-f32.wgsl.jinja", |
| "inputs": { |
| "phase": "\"apply\"", |
| "fusedSoftmax": true, |
| "layout": "\"bsh\"", |
| "kvLayout": "\"bhsd\"", |
| "causalRightAlign": true |
| } |
| }, |
| "bindings": "gqaMatPastApplyFused", |
| "dispatch": { |
| "x": "ceilDiv(gqaHeadDim, gqaMatKeyTile)", |
| "y": "ceilDiv(dim(shapes.queryT, 1), gqaMatQueryTile)", |
| "z": "dim(shapes.queryT, 0) * attrs.num_heads" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "past_kv_bias_materialized_sgmat_f32", |
| "requires": { |
| "features": ["subgroups", "chromium-experimental-subgroup-matrix"], |
| "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] |
| }, |
| "description": "Adds the additive attention bias (batch/head-broadcast at runtime) to the materialized causal scores as they are stored, so the emitted row statistics fold it in for free; everything else matches `past_kv_materialized_sgmat_f32`.", |
| "priority": 36, |
| "when": ["gqaMatPastBiasOk"], |
| "constants": { |
| "qNumHeads": "attrs.num_heads", |
| "kvNumHeads": "attrs.kv_num_heads", |
| "headDim": "gqaHeadDim", |
| "qHidden": "dim(shapes.queryT, 2)", |
| "hasBias": false, |
| "useSubgroups": true, |
| "attentionScaleExpression": "\"select(inverseSqrt(f32(HEAD_DIM)), params.scale, params.scale != 0.0)\"", |
| "materializedSgmatQueryTile": "gqaMatQueryTile", |
| "materializedSgmatKeyTile": "gqaMatKeyTile", |
| "materializedSgmatInnerTile": "gqaMatInnerTile", |
| "materializedSgmatDirectScoreStore": false, |
| "materializedSgmatDirectApplyStore": false, |
| "materializedSgmatRuntimeDirectStore": false, |
| "materializedRowStatsWg": "gqaMatRowStatsWg", |
| "statSlots": "gqaMatPastStatSlots", |
| "statQuerySeq": "dim(shapes.queryT, 1)", |
| "mode": "\"copy\"", |
| "packed": "gqaHeadDim", |
| "kvHeads": "attrs.kv_num_heads", |
| "inputScalar": "gqaScalar", |
| "usesF16": false, |
| "copyWorkgroupSize": "copyWorkgroupSize" |
| }, |
| "intermediates": [ |
| { |
| "id": "materializedScores", |
| "dtype": "float32", |
| "shape": "[dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1) * dim(shapes.presentKeyT, 2)]" |
| }, |
| { "id": "materializedRowStats", "dtype": "float32", "shape": "[gqaMatPastRowStatsElements]" }, |
| { "id": "materializedScorePartials", "dtype": "float32", "shape": "[gqaMatPastScorePartialElements]" } |
| ], |
| "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": "scores", |
| "name": "GroupQueryAttention.PastMaterializedScores", |
| "source": { |
| "shader": "attn-materialized-sgmat-f32.wgsl.jinja", |
| "inputs": { |
| "phase": "\"score\"", |
| "emitRowStats": true, |
| "layout": "\"bsh\"", |
| "kvLayout": "\"bhsd\"", |
| "causalRightAlign": true, |
| "scoreBias": true |
| } |
| }, |
| "bindings": "gqaMatPastBiasScoreStats", |
| "dispatch": { |
| "x": "ceilDiv(dim(shapes.presentKeyT, 2), gqaMatKeyTile)", |
| "y": "ceilDiv(dim(shapes.queryT, 1), gqaMatQueryTile)", |
| "z": "dim(shapes.queryT, 0) * attrs.num_heads" |
| } |
| }, |
| { |
| "id": "rowstats", |
| "name": "GroupQueryAttention.PastMaterializedRowStats", |
| "source": { "shader": "attn-materialized-rowstats-combine-f32.wgsl.jinja", "inputs": {} }, |
| "bindings": "gqaMatPastRowStatsCombine", |
| "dispatch": { |
| "gridStride": "dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)", |
| "workgroupSize": "gqaMatRowStatsWg" |
| } |
| }, |
| { |
| "id": "apply", |
| "name": "GroupQueryAttention.PastMaterializedApply", |
| "source": { |
| "shader": "attn-materialized-sgmat-f32.wgsl.jinja", |
| "inputs": { |
| "phase": "\"apply\"", |
| "fusedSoftmax": true, |
| "layout": "\"bsh\"", |
| "kvLayout": "\"bhsd\"", |
| "causalRightAlign": true |
| } |
| }, |
| "bindings": "gqaMatPastApplyFused", |
| "dispatch": { |
| "x": "ceilDiv(gqaHeadDim, gqaMatKeyTile)", |
| "y": "ceilDiv(dim(shapes.queryT, 1), gqaMatQueryTile)", |
| "z": "dim(shapes.queryT, 0) * attrs.num_heads" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "past_kv_materialized_sgmat_f16", |
| "requires": { |
| "features": ["subgroups", "chromium-experimental-subgroup-matrix"], |
| "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }] |
| }, |
| "description": "Materializes the causal score matrix with `f16` operand tiles feeding `f32`-accumulating subgroup matrices over the shared cache and applies the softmax-normalized weights the same way; scores and row statistics stay `f32`. The score pass skips key tiles past each query tile's causal bound and the apply pass stops its reduction there.", |
| "priority": 35, |
| "when": ["gqaMatPastCoreF16Ok"], |
| "constants": { |
| "qNumHeads": "attrs.num_heads", |
| "kvNumHeads": "attrs.kv_num_heads", |
| "headDim": "gqaHeadDim", |
| "qHidden": "dim(shapes.queryT, 2)", |
| "hasBias": false, |
| "useSubgroups": true, |
| "attentionScaleExpression": "\"select(inverseSqrt(f32(HEAD_DIM)), params.scale, params.scale != 0.0)\"", |
| "materializedSgmatQueryTile": "gqaMatQueryTile", |
| "materializedSgmatKeyTile": "gqaMatKeyTile", |
| "materializedSgmatInnerTile": "gqaMatInnerTile", |
| "materializedSgmatDirectScoreStore": false, |
| "materializedSgmatDirectApplyStore": false, |
| "materializedSgmatRuntimeDirectStore": false, |
| "materializedRowStatsWg": "gqaMatRowStatsWg", |
| "statSlots": "gqaMatPastStatSlots", |
| "statQuerySeq": "dim(shapes.queryT, 1)", |
| "mode": "\"copy\"", |
| "packed": "gqaHeadDim", |
| "kvHeads": "attrs.kv_num_heads", |
| "inputScalar": "gqaScalar", |
| "usesF16": true, |
| "copyWorkgroupSize": "copyWorkgroupSize", |
| "operandF16": true |
| }, |
| "intermediates": [ |
| { |
| "id": "materializedScores", |
| "dtype": "float32", |
| "shape": "[dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1) * dim(shapes.presentKeyT, 2)]" |
| }, |
| { "id": "materializedRowStats", "dtype": "float32", "shape": "[gqaMatPastRowStatsElements]" }, |
| { "id": "materializedScorePartials", "dtype": "float32", "shape": "[gqaMatPastScorePartialElements]" } |
| ], |
| "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": "scores", |
| "name": "GroupQueryAttention.PastMaterializedScores", |
| "source": { |
| "shader": "attn-materialized-sgmat-f32.wgsl.jinja", |
| "inputs": { |
| "phase": "\"score\"", |
| "emitRowStats": true, |
| "layout": "\"bsh\"", |
| "kvLayout": "\"bhsd\"", |
| "causalRightAlign": true |
| } |
| }, |
| "bindings": "gqaMatPastScoreStatsF16", |
| "dispatch": { |
| "x": "ceilDiv(dim(shapes.presentKeyT, 2), gqaMatKeyTile)", |
| "y": "ceilDiv(dim(shapes.queryT, 1), gqaMatQueryTile)", |
| "z": "dim(shapes.queryT, 0) * attrs.num_heads" |
| } |
| }, |
| { |
| "id": "rowstats", |
| "name": "GroupQueryAttention.PastMaterializedRowStats", |
| "source": { "shader": "attn-materialized-rowstats-combine-f32.wgsl.jinja", "inputs": {} }, |
| "bindings": "gqaMatPastRowStatsCombine", |
| "dispatch": { |
| "gridStride": "dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)", |
| "workgroupSize": "gqaMatRowStatsWg" |
| } |
| }, |
| { |
| "id": "apply", |
| "name": "GroupQueryAttention.PastMaterializedApply", |
| "source": { |
| "shader": "attn-materialized-sgmat-f32.wgsl.jinja", |
| "inputs": { |
| "phase": "\"apply\"", |
| "fusedSoftmax": true, |
| "layout": "\"bsh\"", |
| "kvLayout": "\"bhsd\"", |
| "causalRightAlign": true |
| } |
| }, |
| "bindings": "gqaMatPastApplyFusedF16", |
| "dispatch": { |
| "x": "ceilDiv(gqaHeadDim, gqaMatKeyTile)", |
| "y": "ceilDiv(dim(shapes.queryT, 1), gqaMatQueryTile)", |
| "z": "dim(shapes.queryT, 0) * attrs.num_heads" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "past_kv_rotary_materialized_sgmat_f32", |
| "requires": { |
| "features": ["subgroups", "chromium-experimental-subgroup-matrix"], |
| "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] |
| }, |
| "description": "Rotary twin of the causal materialized route: the shared cache already holds rotary-transformed keys, so one preparation pass rotates the query block at its absolute positions and the tile-skipping score/apply passes run unchanged on the prepared buffer.", |
| "priority": 37, |
| "when": ["gqaMatPastRotaryOk"], |
| "constants": { |
| "qNumHeads": "attrs.num_heads", |
| "kvNumHeads": "attrs.kv_num_heads", |
| "headDim": "gqaHeadDim", |
| "qHidden": "dim(shapes.queryT, 2)", |
| "hasBias": false, |
| "useSubgroups": true, |
| "attentionScaleExpression": "\"select(inverseSqrt(f32(HEAD_DIM)), params.scale, params.scale != 0.0)\"", |
| "materializedSgmatQueryTile": "gqaMatQueryTile", |
| "materializedSgmatKeyTile": "gqaMatKeyTile", |
| "materializedSgmatInnerTile": "gqaMatInnerTile", |
| "materializedSgmatDirectScoreStore": false, |
| "materializedSgmatDirectApplyStore": false, |
| "materializedSgmatRuntimeDirectStore": false, |
| "materializedRowStatsWg": "gqaMatRowStatsWg", |
| "statSlots": "gqaMatPastStatSlots", |
| "statQuerySeq": "dim(shapes.queryT, 1)", |
| "mode": "\"copy\"", |
| "packed": "gqaHeadDim", |
| "kvHeads": "attrs.kv_num_heads", |
| "inputScalar": "gqaScalar", |
| "usesF16": false, |
| "copyWorkgroupSize": "copyWorkgroupSize", |
| "half": "gqaHeadDim / 2", |
| "qHeads": "attrs.num_heads", |
| "cosScalar": "\"f16\" if tensorDtypes.cosCacheT == \"float16\" else \"f32\"", |
| "hasQNorm": false |
| }, |
| "intermediates": [ |
| { |
| "id": "qPrep", |
| "dtype": "float32", |
| "shape": "[dim(shapes.queryT, 0), dim(shapes.queryT, 1), dim(shapes.queryT, 2)]" |
| }, |
| { |
| "id": "materializedScores", |
| "dtype": "float32", |
| "shape": "[dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1) * dim(shapes.presentKeyT, 2)]" |
| }, |
| { "id": "materializedRowStats", "dtype": "float32", "shape": "[gqaMatPastRowStatsElements]" }, |
| { "id": "materializedScorePartials", "dtype": "float32", "shape": "[gqaMatPastScorePartialElements]" } |
| ], |
| "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.RotaryQPrep", |
| "shader": "gqa-qprep.wgsl.jinja", |
| "bindings": "rotaryQprep", |
| "dispatch": { |
| "threads": "dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)", |
| "workgroupSize": "copyWorkgroupSize" |
| } |
| }, |
| { |
| "id": "scores", |
| "name": "GroupQueryAttention.PastRotaryMaterializedScores", |
| "source": { |
| "shader": "attn-materialized-sgmat-f32.wgsl.jinja", |
| "inputs": { |
| "phase": "\"score\"", |
| "emitRowStats": true, |
| "layout": "\"bsh\"", |
| "kvLayout": "\"bhsd\"", |
| "causalRightAlign": true |
| } |
| }, |
| "bindings": "gqaMatPastRotaryScoreStats", |
| "dispatch": { |
| "x": "ceilDiv(dim(shapes.presentKeyT, 2), gqaMatKeyTile)", |
| "y": "ceilDiv(dim(shapes.queryT, 1), gqaMatQueryTile)", |
| "z": "dim(shapes.queryT, 0) * attrs.num_heads" |
| } |
| }, |
| { |
| "id": "rowstats", |
| "name": "GroupQueryAttention.PastRotaryMaterializedRowStats", |
| "source": { "shader": "attn-materialized-rowstats-combine-f32.wgsl.jinja", "inputs": {} }, |
| "bindings": "gqaMatPastRowStatsCombine", |
| "dispatch": { |
| "gridStride": "dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)", |
| "workgroupSize": "gqaMatRowStatsWg" |
| } |
| }, |
| { |
| "id": "apply", |
| "name": "GroupQueryAttention.PastRotaryMaterializedApply", |
| "source": { |
| "shader": "attn-materialized-sgmat-f32.wgsl.jinja", |
| "inputs": { |
| "phase": "\"apply\"", |
| "fusedSoftmax": true, |
| "layout": "\"bsh\"", |
| "kvLayout": "\"bhsd\"", |
| "causalRightAlign": true |
| } |
| }, |
| "bindings": "gqaMatPastApplyFused", |
| "dispatch": { |
| "x": "ceilDiv(gqaHeadDim, gqaMatKeyTile)", |
| "y": "ceilDiv(dim(shapes.queryT, 1), gqaMatQueryTile)", |
| "z": "dim(shapes.queryT, 0) * attrs.num_heads" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "past_kv_rotary_materialized_sgmat_f16", |
| "requires": { |
| "features": ["subgroups", "chromium-experimental-subgroup-matrix"], |
| "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }] |
| }, |
| "description": "Materializes the causal score matrix with `f16` operand tiles feeding `f32`-accumulating subgroup matrices over the shared cache and applies the softmax-normalized weights the same way; scores and row statistics stay `f32`. The score pass skips key tiles past each query tile's causal bound and the apply pass stops its reduction there.", |
| "priority": 37, |
| "when": ["gqaMatPastRotaryF16Ok"], |
| "constants": { |
| "qNumHeads": "attrs.num_heads", |
| "kvNumHeads": "attrs.kv_num_heads", |
| "headDim": "gqaHeadDim", |
| "qHidden": "dim(shapes.queryT, 2)", |
| "hasBias": false, |
| "useSubgroups": true, |
| "attentionScaleExpression": "\"select(inverseSqrt(f32(HEAD_DIM)), params.scale, params.scale != 0.0)\"", |
| "materializedSgmatQueryTile": "gqaMatQueryTile", |
| "materializedSgmatKeyTile": "gqaMatKeyTile", |
| "materializedSgmatInnerTile": "gqaMatInnerTile", |
| "materializedSgmatDirectScoreStore": false, |
| "materializedSgmatDirectApplyStore": false, |
| "materializedSgmatRuntimeDirectStore": false, |
| "materializedRowStatsWg": "gqaMatRowStatsWg", |
| "statSlots": "gqaMatPastStatSlots", |
| "statQuerySeq": "dim(shapes.queryT, 1)", |
| "mode": "\"copy\"", |
| "packed": "gqaHeadDim", |
| "kvHeads": "attrs.kv_num_heads", |
| "inputScalar": "gqaScalar", |
| "usesF16": true, |
| "copyWorkgroupSize": "copyWorkgroupSize", |
| "half": "gqaHeadDim / 2", |
| "qHeads": "attrs.num_heads", |
| "cosScalar": "\"f16\" if tensorDtypes.cosCacheT == \"float16\" else \"f32\"", |
| "hasQNorm": false, |
| "operandF16": true |
| }, |
| "intermediates": [ |
| { |
| "id": "qPrep", |
| "dtype": "float16", |
| "shape": "[dim(shapes.queryT, 0), dim(shapes.queryT, 1), dim(shapes.queryT, 2)]" |
| }, |
| { |
| "id": "materializedScores", |
| "dtype": "float32", |
| "shape": "[dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1) * dim(shapes.presentKeyT, 2)]" |
| }, |
| { "id": "materializedRowStats", "dtype": "float32", "shape": "[gqaMatPastRowStatsElements]" }, |
| { "id": "materializedScorePartials", "dtype": "float32", "shape": "[gqaMatPastScorePartialElements]" } |
| ], |
| "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.RotaryQPrep", |
| "shader": "gqa-qprep.wgsl.jinja", |
| "bindings": "rotaryQprepF16", |
| "dispatch": { |
| "threads": "dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)", |
| "workgroupSize": "copyWorkgroupSize" |
| }, |
| "constants": { "qPrepF16": true } |
| }, |
| { |
| "id": "scores", |
| "name": "GroupQueryAttention.PastRotaryMaterializedScores", |
| "source": { |
| "shader": "attn-materialized-sgmat-f32.wgsl.jinja", |
| "inputs": { |
| "phase": "\"score\"", |
| "emitRowStats": true, |
| "layout": "\"bsh\"", |
| "kvLayout": "\"bhsd\"", |
| "causalRightAlign": true |
| } |
| }, |
| "bindings": "gqaMatPastRotaryScoreStatsF16", |
| "dispatch": { |
| "x": "ceilDiv(dim(shapes.presentKeyT, 2), gqaMatKeyTile)", |
| "y": "ceilDiv(dim(shapes.queryT, 1), gqaMatQueryTile)", |
| "z": "dim(shapes.queryT, 0) * attrs.num_heads" |
| } |
| }, |
| { |
| "id": "rowstats", |
| "name": "GroupQueryAttention.PastRotaryMaterializedRowStats", |
| "source": { "shader": "attn-materialized-rowstats-combine-f32.wgsl.jinja", "inputs": {} }, |
| "bindings": "gqaMatPastRowStatsCombine", |
| "dispatch": { |
| "gridStride": "dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)", |
| "workgroupSize": "gqaMatRowStatsWg" |
| } |
| }, |
| { |
| "id": "apply", |
| "name": "GroupQueryAttention.PastRotaryMaterializedApply", |
| "source": { |
| "shader": "attn-materialized-sgmat-f32.wgsl.jinja", |
| "inputs": { |
| "phase": "\"apply\"", |
| "fusedSoftmax": true, |
| "layout": "\"bsh\"", |
| "kvLayout": "\"bhsd\"", |
| "causalRightAlign": true |
| } |
| }, |
| "bindings": "gqaMatPastApplyFusedF16", |
| "dispatch": { |
| "x": "ceilDiv(gqaHeadDim, gqaMatKeyTile)", |
| "y": "ceilDiv(dim(shapes.queryT, 1), gqaMatQueryTile)", |
| "z": "dim(shapes.queryT, 0) * attrs.num_heads" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "new_kv_past_materialized_sgmat_f32", |
| "requires": { |
| "features": ["subgroups", "chromium-experimental-subgroup-matrix"], |
| "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] |
| }, |
| "description": "Append-route twin of the causal materialized route for chunked prefill: one merge pass concatenates the past cache with the new key/value rows, and the tile-skipping score/apply passes run over the merged cache with the same right-aligned causal bound.", |
| "priority": 35, |
| "when": ["gqaMatNewPastOk"], |
| "constants": { |
| "qNumHeads": "attrs.num_heads", |
| "kvNumHeads": "attrs.kv_num_heads", |
| "headDim": "gqaHeadDim", |
| "qHidden": "dim(shapes.queryT, 2)", |
| "hasBias": false, |
| "useSubgroups": true, |
| "attentionScaleExpression": "\"select(inverseSqrt(f32(HEAD_DIM)), params.scale, params.scale != 0.0)\"", |
| "materializedSgmatQueryTile": "gqaMatQueryTile", |
| "materializedSgmatKeyTile": "gqaMatKeyTile", |
| "materializedSgmatInnerTile": "gqaMatInnerTile", |
| "materializedSgmatDirectScoreStore": false, |
| "materializedSgmatDirectApplyStore": false, |
| "materializedSgmatRuntimeDirectStore": false, |
| "materializedRowStatsWg": "gqaMatRowStatsWg", |
| "statSlots": "gqaMatPastStatSlots", |
| "statQuerySeq": "dim(shapes.queryT, 1)", |
| "mode": "\"merge\"", |
| "packed": "gqaHeadDim", |
| "kvHeads": "attrs.kv_num_heads", |
| "inputScalar": "gqaScalar", |
| "usesF16": false, |
| "copyWorkgroupSize": "copyWorkgroupSize" |
| }, |
| "intermediates": [ |
| { |
| "id": "materializedScores", |
| "dtype": "float32", |
| "shape": "[dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1) * dim(shapes.presentKeyT, 2)]" |
| }, |
| { "id": "materializedRowStats", "dtype": "float32", "shape": "[gqaMatPastRowStatsElements]" }, |
| { "id": "materializedScorePartials", "dtype": "float32", "shape": "[gqaMatPastScorePartialElements]" } |
| ], |
| "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": "scores", |
| "name": "GroupQueryAttention.NewPastMaterializedScores", |
| "source": { |
| "shader": "attn-materialized-sgmat-f32.wgsl.jinja", |
| "inputs": { |
| "phase": "\"score\"", |
| "emitRowStats": true, |
| "layout": "\"bsh\"", |
| "kvLayout": "\"bhsd\"", |
| "causalRightAlign": true |
| } |
| }, |
| "bindings": "gqaMatPastScoreStats", |
| "dispatch": { |
| "x": "ceilDiv(dim(shapes.presentKeyT, 2), gqaMatKeyTile)", |
| "y": "ceilDiv(dim(shapes.queryT, 1), gqaMatQueryTile)", |
| "z": "dim(shapes.queryT, 0) * attrs.num_heads" |
| } |
| }, |
| { |
| "id": "rowstats", |
| "name": "GroupQueryAttention.NewPastMaterializedRowStats", |
| "source": { "shader": "attn-materialized-rowstats-combine-f32.wgsl.jinja", "inputs": {} }, |
| "bindings": "gqaMatPastRowStatsCombine", |
| "dispatch": { |
| "gridStride": "dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)", |
| "workgroupSize": "gqaMatRowStatsWg" |
| } |
| }, |
| { |
| "id": "apply", |
| "name": "GroupQueryAttention.NewPastMaterializedApply", |
| "source": { |
| "shader": "attn-materialized-sgmat-f32.wgsl.jinja", |
| "inputs": { |
| "phase": "\"apply\"", |
| "fusedSoftmax": true, |
| "layout": "\"bsh\"", |
| "kvLayout": "\"bhsd\"", |
| "causalRightAlign": true |
| } |
| }, |
| "bindings": "gqaMatPastApplyFused", |
| "dispatch": { |
| "x": "ceilDiv(gqaHeadDim, gqaMatKeyTile)", |
| "y": "ceilDiv(dim(shapes.queryT, 1), gqaMatQueryTile)", |
| "z": "dim(shapes.queryT, 0) * attrs.num_heads" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "window_shift_materialized_sgmat_f32", |
| "requires": { |
| "features": ["subgroups", "chromium-experimental-subgroup-matrix"], |
| "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] |
| }, |
| "description": "Windowed-cache twin of the causal materialized route for chunked prefill: the shift pass compacts surviving cache rows and appends the chunk, the score pass masks each row's sliding-window floor alongside the causal bound so the row statistics fold both out, and the apply pass starts each reduction at the first key tile the window can reach.", |
| "priority": 36, |
| "when": ["gqaMatWindowOk"], |
| "constants": { |
| "qNumHeads": "attrs.num_heads", |
| "kvNumHeads": "attrs.kv_num_heads", |
| "headDim": "gqaHeadDim", |
| "qHidden": "dim(shapes.queryT, 2)", |
| "hasBias": false, |
| "useSubgroups": true, |
| "attentionScaleExpression": "\"select(inverseSqrt(f32(HEAD_DIM)), params.scale, params.scale != 0.0)\"", |
| "materializedSgmatQueryTile": "gqaMatQueryTile", |
| "materializedSgmatKeyTile": "gqaMatKeyTile", |
| "materializedSgmatInnerTile": "gqaMatInnerTile", |
| "materializedSgmatDirectScoreStore": false, |
| "materializedSgmatDirectApplyStore": false, |
| "materializedSgmatRuntimeDirectStore": false, |
| "materializedRowStatsWg": "gqaMatRowStatsWg", |
| "statSlots": "gqaMatPastStatSlots", |
| "statQuerySeq": "dim(shapes.queryT, 1)", |
| "mode": "\"window_shift\"", |
| "packed": "gqaHeadDim", |
| "kvHeads": "attrs.kv_num_heads", |
| "inputScalar": "gqaScalar", |
| "zeroScalar": "gqaScalar", |
| "usesF16": false, |
| "copyWorkgroupSize": "copyWorkgroupSize" |
| }, |
| "intermediates": [ |
| { |
| "id": "materializedScores", |
| "dtype": "float32", |
| "shape": "[dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1) * dim(shapes.presentKeyT, 2)]" |
| }, |
| { "id": "materializedRowStats", "dtype": "float32", "shape": "[gqaMatPastRowStatsElements]" }, |
| { "id": "materializedScorePartials", "dtype": "float32", "shape": "[gqaMatPastScorePartialElements]" } |
| ], |
| "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": "scores", |
| "name": "GroupQueryAttention.WindowMaterializedScores", |
| "source": { |
| "shader": "attn-materialized-sgmat-f32.wgsl.jinja", |
| "inputs": { |
| "phase": "\"score\"", |
| "emitRowStats": true, |
| "layout": "\"bsh\"", |
| "kvLayout": "\"bhsd\"", |
| "causalRightAlign": true, |
| "scoreWindow": true, |
| "useSeqlens": true |
| } |
| }, |
| "bindings": "gqaMatWindowScoreStats", |
| "dispatch": { |
| "x": "ceilDiv(dim(shapes.presentKeyT, 2), gqaMatKeyTile)", |
| "y": "ceilDiv(dim(shapes.queryT, 1), gqaMatQueryTile)", |
| "z": "dim(shapes.queryT, 0) * attrs.num_heads" |
| } |
| }, |
| { |
| "id": "rowstats", |
| "name": "GroupQueryAttention.WindowMaterializedRowStats", |
| "source": { "shader": "attn-materialized-rowstats-combine-f32.wgsl.jinja", "inputs": {} }, |
| "bindings": "gqaMatPastRowStatsCombine", |
| "dispatch": { |
| "gridStride": "dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)", |
| "workgroupSize": "gqaMatRowStatsWg" |
| } |
| }, |
| { |
| "id": "apply", |
| "name": "GroupQueryAttention.WindowMaterializedApply", |
| "source": { |
| "shader": "attn-materialized-sgmat-f32.wgsl.jinja", |
| "inputs": { |
| "phase": "\"apply\"", |
| "fusedSoftmax": true, |
| "layout": "\"bsh\"", |
| "kvLayout": "\"bhsd\"", |
| "causalRightAlign": true, |
| "scoreWindow": true, |
| "useSeqlens": true |
| } |
| }, |
| "bindings": "gqaMatWindowApplyFused", |
| "dispatch": { |
| "x": "ceilDiv(gqaHeadDim, gqaMatKeyTile)", |
| "y": "ceilDiv(dim(shapes.queryT, 1), gqaMatQueryTile)", |
| "z": "dim(shapes.queryT, 0) * attrs.num_heads" |
| } |
| } |
| ] |
| }, |
| { |
| "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<f32>\"", |
| "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<f32>\"", |
| "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<f16>\"", |
| "outputScalar": "\"f16\"", |
| "presentScalar": "qkvCacheScalar", |
| "inputElement": "\"vec4<f16>\"", |
| "outputElement": "\"vec4<f16>\"", |
| "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<f16>\"", |
| "outputScalar": "\"f16\"", |
| "presentScalar": "qkvCacheScalar", |
| "inputElement": "\"vec4<f16>\"", |
| "outputElement": "\"vec4<f16>\"", |
| "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<i32>\"", |
| "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<i32>\"", |
| "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)" |
| } |
| } |
| ] |
| } |
| ] |
| } |
|
|
