diff --git "a/build/webgpu/manifest.json" "b/build/webgpu/manifest.json" --- "a/build/webgpu/manifest.json" +++ "b/build/webgpu/manifest.json" @@ -2,155 +2,42 @@ "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" + "inputs": { + "queryT": { "onnx": "query", "dtype": "T", "rank": 3 }, + "keyT": { "onnx": "key", "dtype": "T", "rank": 3 }, + "valueT": { "onnx": "value", "dtype": "T", "rank": 3 }, + "pastKeyT": { "onnx": "past_key", "dtype": "T_CACHE", "rank": 4, "optional": true }, + "pastValueT": { "onnx": "past_value", "dtype": "T_CACHE", "rank": 4, "optional": true }, + "seqlensKT": { "onnx": "seqlens_k", "dtype": "M", "rank": 1 }, + "totalSequenceLengthT": { "onnx": "total_sequence_length", "dtype": "M", "rank": 1 }, + "cosCacheT": { "onnx": "cos_cache", "dtype": "T", "rank": 2, "optional": true }, + "sinCacheT": { "onnx": "sin_cache", "dtype": "T", "rank": 2, "optional": true }, + "attentionBiasT": { "onnx": "attention_bias", "dtype": "T", "rank": 4, "optional": true }, + "headSinkT": { "onnx": "head_sink", "dtype": "T", "rank": 1, "optional": true }, + "kScaleT": { "onnx": "k_scale", "dtype": "T_KV_SCALE", "rank": 1, "optional": true }, + "vScaleT": { "onnx": "v_scale", "dtype": "T_KV_SCALE", "rank": 1, "optional": true }, + "qNormWeightT": { "onnx": "q_norm_weight", "dtype": "T", "rank": 1, "optional": true }, + "kNormWeightT": { "onnx": "k_norm_weight", "dtype": "T", "rank": 1, "optional": true } }, - "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`." + "outputs": { + "outputT": { "onnx": "output", "dtype": "T", "rank": 3, "shape": "shapes.queryT" }, + "presentKeyT": { "onnx": "present_key", "dtype": "T_CACHE", "rank": 4 }, + "presentValueT": { "onnx": "present_value", "dtype": "T_CACHE", "rank": 4 } + }, + "attributes": { + "causal": { "default": 1 }, + "do_rotary": { "default": 0 }, + "k_quant_type": { "default": "NONE" }, + "local_window_size": { "default": -1 }, + "qk_norm_epsilon": { "default": 0.000001 }, + "sliding_window_cache": { "default": 0 }, + "smooth_softmax": { "default": -1 }, + "softcap": { "default": 0 }, + "v_quant_type": { "default": "NONE" }, + "kv_cache_bit_width": {}, + "kv_num_heads": {}, + "num_heads": {}, + "scale": {} }, "attributeConstraints": { "causal": { "values": [0, 1] }, @@ -168,56 +55,37 @@ "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 + "WORKGROUP_SIZE": { "default": 256 }, + "COPY_WORKGROUP_SIZE": { "default": 64 }, + "SCALAR_WORKGROUP_SIZE": { "default": 64 }, + "COOPERATIVE_QUERY_THRESHOLD": { "default": 6144 }, + "MAX_SPLITS": { "default": 16 }, + "QKV_SPLIT_TILE_K": { "default": 128 }, + "CACHED_SPLIT_TILE_K": { "default": 256 }, + "FLASH_MIN_HEAD_DIM": { "default": 32 }, + "FLASH_MAX_HEAD_DIM": { "default": 256 }, + "PREFILL_QUERY_TILE": { "default": 32 }, + "PREFILL_LANES_PER_QUERY": { "default": 8 }, + "PREFILL_QUERY_TILE_F16": { "default": 16 }, + "PREFILL_LANES_PER_QUERY_F16": { "default": 8 }, + "QKV_PREFILL_MIN_QUERY_TOKENS": { "default": 31 }, + "QKV_DECODE_MIN_KV_TOKENS": { "default": 512 }, + "CACHED_FLASH_MIN_HEAD_DIM": { "default": 16 }, + "CACHED_DECODE_MIN_KV_TOKENS": { "default": 1024 }, + "CLUSTER_TILE_K_F32": { "default": 8 }, + "CLUSTER_TILE_K_F16": { "default": 8 }, + "NO_SG_TILE_K_MAX": { "default": 16 }, + "COOPERATIVE_WORKGROUP_SIZE": { "default": 32 }, + "CLUSTER_MAX_SLICE": { "default": 8 }, + "MATERIALIZED_SGMAT_QUERY_TILE": { "default": 64 }, + "MATERIALIZED_SGMAT_KEY_TILE": { "default": 64 }, + "MATERIALIZED_SGMAT_INNER_TILE": { "default": 32 }, + "MATERIALIZED_SOFTMAX_WORKGROUP_SIZE": { "default": 256 }, + "MATERIALIZED_SGMAT_MIN_SEQ": { "default": 256 }, + "MATERIALIZED_SGMAT_PAST_MIN_SEQ": { "default": 256 }, + "MATERIALIZED_SGMAT_WINDOW_MIN_QSEQ": { "default": 128 }, + "MATERIALIZED_SGMAT_APPEND_MIN_QSEQ": { "default": 128 } }, "derive": { "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)", @@ -228,7 +96,7 @@ "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", + "headDimPlan": "dim(shapes.queryT, 2) / attrs.num_heads if (ranks.queryT == 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", @@ -241,28 +109,28 @@ "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)", + "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.queryT == 3 and ranks.keyT == 3 and ranks.valueT == 3 and ranks.outputT == 3 and ranks.presentKeyT == 4 and ranks.presentValueT == 4 and (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 == \"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.queryT, 2) % attrs.num_heads == 0 and dim(shapes.keyT, 2) == attrs.kv_num_heads * headDimPlan and dim(shapes.valueT, 2) == dim(shapes.keyT, 2) and dim(shapes.queryT, 0) == dim(shapes.keyT, 0) and dim(shapes.queryT, 0) == dim(shapes.valueT, 0) and dim(shapes.keyT, 1) == dim(shapes.valueT, 1) and dim(shapes.outputT, 0) == dim(shapes.queryT, 0) and dim(shapes.outputT, 1) == dim(shapes.queryT, 1) and dim(shapes.outputT, 2) == dim(shapes.queryT, 2) and dim(shapes.presentKeyT, 0) == dim(shapes.queryT, 0) and dim(shapes.presentKeyT, 1) == attrs.kv_num_heads and dim(shapes.presentKeyT, 2) == dim(shapes.keyT, 1) and dim(shapes.presentKeyT, 3) == headDimPlan 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", + "flashShapeOk": "qkvContractOk and flashWorkgroupOk and headDimPlan % 4 == 0 and headDimPlan >= tunables.FLASH_MIN_HEAD_DIM and headDimPlan <= tunables.FLASH_MAX_HEAD_DIM and (dim(shapes.queryT, 1) * attrs.num_heads >= 248 or (dim(shapes.queryT, 1) == 1 and dim(shapes.keyT, 1) >= tunables.QKV_DECODE_MIN_KV_TOKENS) or (dim(shapes.queryT, 1) > 1 and dim(shapes.keyT, 1) >= 2048)) and dim(shapes.queryT, 0) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and dim(shapes.queryT, 1) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and attrs.num_heads <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", + "decodeSplitKShapeOk": "qkvContractOk and flashWorkgroupOk and dim(shapes.queryT, 1) == 1 and dim(shapes.keyT, 1) >= tunables.QKV_DECODE_MIN_KV_TOKENS and headDimPlan % 4 == 0 and headDimPlan >= tunables.FLASH_MIN_HEAD_DIM and headDimPlan <= tunables.FLASH_MAX_HEAD_DIM and dim(shapes.queryT, 0) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and attrs.num_heads <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", + "fallbackShapeOk": "qkvContractOk and headDimPlan <= deviceWorkgroupCap and headDimPlan <= tunables.FLASH_MAX_HEAD_DIM", "gqaHeadDim": "dim(shapes.queryT, 2) / attrs.num_heads", - "qkvScalar": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"", - "qkvInputVec4": "\"vec4\" if tensorDtypes.query == \"float16\" else \"vec4\"", + "qkvScalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"", + "qkvInputVec4": "\"vec4\" if tensorDtypes.queryT == \"float16\" else \"vec4\"", "qkvCacheScalar": "\"f16\" if tensorDtypes.presentKeyT == \"float16\" else \"f32\"", "qkvCacheVec4": "\"vec4\" if tensorDtypes.presentKeyT == \"float16\" else \"vec4\"", - "qkvUsesF16": "tensorDtypes.query == \"float16\" or tensorDtypes.presentKeyT == \"float16\"", + "qkvUsesF16": "tensorDtypes.queryT == \"float16\" or tensorDtypes.presentKeyT == \"float16\"", "gqaScalar": "\"f16\" if tensorDtypes.queryT == \"float16\" else \"f32\"", "gqaInputVec4": "\"vec4\" if tensorDtypes.queryT == \"float16\" else \"vec4\"", - "qkvNumSplits": "min(tunables.MAX_SPLITS, ceilDiv(dim(shapes.key, 1), tunables.QKV_SPLIT_TILE_K))", + "qkvNumSplits": "min(tunables.MAX_SPLITS, ceilDiv(dim(shapes.keyT, 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))", + "qkvPrefillQueryTile": "tunables.PREFILL_QUERY_TILE_F16 if tensorDtypes.queryT == \"float16\" else tunables.PREFILL_QUERY_TILE", + "qkvPrefillLanesPerQuery": "max(tunables.PREFILL_LANES_PER_QUERY_F16 if tensorDtypes.queryT == \"float16\" else tunables.PREFILL_LANES_PER_QUERY, ceilDiv(headDimPlan / 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", @@ -274,7 +142,7 @@ "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", + "gqaDispatchFits": "dim(shapes.queryT, 0) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and attrs.num_heads <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "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", @@ -285,30 +153,30 @@ "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", + "prefillFlashShapeOk": "cachedPrefillClusterWorkgroupOk and gqaHeadDim % (4 * cachedPrefillLanesPerQuery) == 0 and gqaFlashHd and dim(shapes.queryT, 1) >= cachedPrefillQueryTile and gqaDispatchFits and ceilDiv(dim(shapes.queryT, 1), cachedPrefillQueryTile) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "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", + "qkvSplitScratchBytes": "dim(shapes.queryT, 0) * attrs.num_heads * qkvNumSplits * headDimPlan * 4", + "qkvSplitStatsBytes": "2 * dim(shapes.queryT, 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", + "qkvTiledWorkgroupBytes": "dim(shapes.valueT, 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)", + "qkvClusterTileK": "tunables.CLUSTER_TILE_K_F32 if tensorDtypes.queryT == \"float32\" and headDimPlan <= tunables.QKV_SPLIT_TILE_K else tunables.CLUSTER_TILE_K_F16", + "qkvClusterWorkgroupBytes": "qkvClusterTileK * headDimPlan * (8 if tensorDtypes.queryT == \"float32\" else 4)", "qkvClusterStorageOk": "qkvClusterWorkgroupBytes <= device.limits.maxComputeWorkgroupStorageSize", - "qkvF32ClusterRegisterGeometry": "tensorDtypes.query != \"float32\" or (headDim % 4 == 0 and headDim / (4 * qkvPrefillLanesPerQuery) <= tunables.CLUSTER_MAX_SLICE)", + "qkvF32ClusterRegisterGeometry": "tensorDtypes.queryT != \"float32\" or (headDimPlan % 4 == 0 and headDimPlan / (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", + "qkvNoSgTileBytes": "headDimPlan * (8 if tensorDtypes.queryT == \"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", @@ -330,15 +198,15 @@ "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", + "gqaMatScoreBytes": "dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1) * dim(shapes.keyT, 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", + "gqaMatBuffersFit": "numel(shapes.queryT) * 4 <= device.limits.maxStorageBufferBindingSize and numel(shapes.keyT) * 4 <= device.limits.maxStorageBufferBindingSize and numel(shapes.valueT) * 4 <= device.limits.maxStorageBufferBindingSize", + "gqaMatDispatchFits": "ceilDiv(dim(shapes.keyT, 1), gqaMatKeyTile) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(dim(shapes.queryT, 1), gqaMatQueryTile) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and dim(shapes.queryT, 0) * attrs.num_heads <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", + "gqaMatStatSlots": "ceilDiv(dim(shapes.keyT, 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", + "gqaMatRowStatsElements": "dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1) * 2", + "gqaMatScorePartialElements": "dim(shapes.queryT, 0) * attrs.num_heads * gqaMatStatSlots * dim(shapes.queryT, 1) * 2", + "gqaMatCoreOk": "qkvContractOk and attrs.causal == 0 and tensorDtypes.queryT == \"float32\" and attrs.local_window_size == -1 and headDimPlan >= 64 and headDimPlan <= 256 and headDimPlan % 16 == 0 and dim(shapes.queryT, 1) >= tunables.MATERIALIZED_SGMAT_MIN_SEQ and dim(shapes.keyT, 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", @@ -347,7 +215,8 @@ "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", + "gqaMatPastDispatchFits": "ceilDiv(gqaMatPastKvSeq, gqaMatKeyTile) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(gqaMatPastQSeq, gqaMatQueryTile) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and dim(shapes.queryT, 0) * attrs.num_heads <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", + "gqaMatShareAppendOk": "shareAppendOk and plainAttentionOptions and tensorDtypes.queryT == \"float32\" and attrs.local_window_size == -1 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_APPEND_MIN_QSEQ and dim(shapes.presentKeyT, 2) >= dim(shapes.queryT, 1) and device.features.has(\"subgroups\") and wave32Effective and device.features.has(\"chromium-experimental-subgroup-matrix\") and gqaMatPastScoreFits and gqaMatPastBuffersFit and gqaMatResourcesFit and gqaMatPastDispatchFits", "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", @@ -356,2025 +225,454 @@ "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" }, - { - "name": "partial_stats", - "semantic": "partialStats", - "buffer": { "type": "storage" }, - "elementType": "vec2" - }, - { - "name": "params", - "semantic": "kernel.params", - "buffer": { "type": "uniform" }, - "struct": { - "name": "Params", - "fields": [ - { "name": "kvSeq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" }, - { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, - { - "name": "windowSize", - "type": "u32", - "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" - } - ] - } - } - ], - "splitSeqlensAttention": [ - { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputVec4" }, - { "name": "key", "arg": "presentKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputVec4" }, - { - "name": "value", - "arg": "presentValueT", - "buffer": { "type": "read-only-storage" }, - "elementType": "$inputVec4" - }, - { "name": "partial_out", "semantic": "partialOut", "buffer": { "type": "storage" }, "elementType": "vec4" }, - { - "name": "partial_stats", - "semantic": "partialStats", - "buffer": { "type": "storage" }, - "elementType": "vec2" - }, - { "name": "seqlens_k", "arg": "seqlensKT", "buffer": { "type": "read-only-storage" }, "elementType": "i32" }, - { - "name": "params", - "semantic": "kernel.params", - "buffer": { "type": "uniform" }, - "struct": { - "name": "Params", - "fields": [ - { "name": "kvSeq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" }, - { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, - { - "name": "windowSize", - "type": "u32", - "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" - } - ] - } - } - ], - "splitMerge": [ - { - "name": "partial_out", - "semantic": "partialOut", - "buffer": { "type": "read-only-storage" }, - "elementType": "vec4" - }, - { - "name": "partial_stats", - "semantic": "partialStats", - "buffer": { "type": "read-only-storage" }, - "elementType": "vec2" - }, - { "name": "output", "arg": "outputT", "buffer": { "type": "storage" }, "elementType": "$inputVec4" }, - { - "name": "params", - "semantic": "kernel.params", - "buffer": { "type": "uniform" }, - "struct": { - "name": "Params", - "fields": [{ "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }] - } - } - ], - "splitMergeNoParams": [ - { - "name": "partial_out", - "semantic": "partialOut", - "buffer": { "type": "read-only-storage" }, - "elementType": "vec4" - }, - { - "name": "partial_stats", - "semantic": "partialStats", - "buffer": { "type": "read-only-storage" }, - "elementType": "vec2" - }, - { "name": "output", "arg": "outputT", "buffer": { "type": "storage" }, "elementType": "$inputVec4" } - ], - "flashPrefillPastSoftcap": [ - { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputVec4" }, - { "name": "key", "arg": "presentKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputVec4" }, - { - "name": "value", - "arg": "presentValueT", - "buffer": { "type": "read-only-storage" }, - "elementType": "$inputVec4" - }, - { "name": "output", "arg": "outputT", "buffer": { "type": "storage" }, "elementType": "$inputVec4" }, - { - "name": "params", - "semantic": "kernel.params", - "buffer": { "type": "uniform" }, - "struct": { - "name": "Params", - "fields": [ - { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, - { "name": "kvSeq", "type": "u32", "value": "dim(shapes.pastKeyT, 2)" }, - { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, - { "name": "isCausal", "type": "u32", "value": 1 }, - { - "name": "windowSize", - "type": "u32", - "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" - }, - { "name": "softcap", "type": "f32", "value": "attrs.softcap" } - ] - } - } - ], - "flashPrefillPastHeadSink": [ - { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputVec4" }, - { "name": "key", "arg": "presentKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputVec4" }, - { - "name": "value", - "arg": "presentValueT", - "buffer": { "type": "read-only-storage" }, - "elementType": "$inputVec4" - }, - { - "name": "head_sink", - "arg": "headSinkT", - "buffer": { "type": "read-only-storage" }, - "elementType": "$inputScalar" - }, - { "name": "output", "arg": "outputT", "buffer": { "type": "storage" }, "elementType": "$inputVec4" }, - { - "name": "params", - "semantic": "kernel.params", - "buffer": { "type": "uniform" }, - "struct": { - "name": "Params", - "fields": [ - { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, - { "name": "kvSeq", "type": "u32", "value": "dim(shapes.pastKeyT, 2)" }, - { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, - { "name": "isCausal", "type": "u32", "value": 1 }, - { - "name": "windowSize", - "type": "u32", - "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" - } - ] - } - } - ], - "flashPrefillBiasHeadSink": [ - { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputVec4" }, - { "name": "key", "arg": "presentKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputVec4" }, - { - "name": "value", - "arg": "presentValueT", - "buffer": { "type": "read-only-storage" }, - "elementType": "$inputVec4" - }, - { - "name": "attn_mask", - "arg": "attentionBiasT", - "buffer": { "type": "read-only-storage" }, - "elementType": "$inputScalar" - }, - { - "name": "head_sink", - "arg": "headSinkT", - "buffer": { "type": "read-only-storage" }, - "elementType": "$inputScalar" - }, - { "name": "output", "arg": "outputT", "buffer": { "type": "storage" }, "elementType": "$inputVec4" }, - { - "name": "params", - "semantic": "kernel.params", - "buffer": { "type": "uniform" }, - "struct": { - "name": "Params", - "fields": [ - { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, - { "name": "kvSeq", "type": "u32", "value": "dim(shapes.pastKeyT, 2)" }, - { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, - { "name": "isCausal", "type": "u32", "value": 1 }, - { - "name": "windowSize", - "type": "u32", - "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" - }, - { - "name": "maskBatchStride", - "type": "u32", - "value": "0 if dim(shapes.attentionBiasT, 0) == 1 else dim(shapes.attentionBiasT, 1) * dim(shapes.attentionBiasT, 2) * dim(shapes.attentionBiasT, 3)" - }, - { - "name": "maskHeadStride", - "type": "u32", - "value": "0 if dim(shapes.attentionBiasT, 1) == 1 else dim(shapes.attentionBiasT, 2) * dim(shapes.attentionBiasT, 3)" - }, - { "name": "maskSeqStride", "type": "u32", "value": "dim(shapes.attentionBiasT, 3)" } - ] - } - } - ], - "flashPrefillPast": [ - { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputVec4" }, - { "name": "key", "arg": "presentKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputVec4" }, - { - "name": "value", - "arg": "presentValueT", - "buffer": { "type": "read-only-storage" }, - "elementType": "$inputVec4" - }, - { "name": "output", "arg": "outputT", "buffer": { "type": "storage" }, "elementType": "$inputVec4" }, - { - "name": "params", - "semantic": "kernel.params", - "buffer": { "type": "uniform" }, - "struct": { - "name": "Params", - "fields": [ - { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, - { "name": "kvSeq", "type": "u32", "value": "dim(shapes.pastKeyT, 2)" }, - { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, - { "name": "isCausal", "type": "u32", "value": 1 }, - { - "name": "windowSize", - "type": "u32", - "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" - } - ] - } - } - ], - "flashPrefillSeqlens": [ - { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputVec4" }, - { "name": "key", "arg": "presentKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputVec4" }, - { - "name": "value", - "arg": "presentValueT", - "buffer": { "type": "read-only-storage" }, - "elementType": "$inputVec4" - }, - { "name": "output", "arg": "outputT", "buffer": { "type": "storage" }, "elementType": "$inputVec4" }, - { "name": "seqlens_k", "arg": "seqlensKT", "buffer": { "type": "read-only-storage" }, "elementType": "i32" }, - { - "name": "params", - "semantic": "kernel.params", - "buffer": { "type": "uniform" }, - "struct": { - "name": "Params", - "fields": [ - { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, - { "name": "kvSeq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" }, - { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, - { "name": "isCausal", "type": "u32", "value": 1 }, - { - "name": "windowSize", - "type": "u32", - "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" - } - ] - } - } - ], - "flashPrefillNew": [ - { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputVec4" }, - { "name": "key", "arg": "presentKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputVec4" }, - { - "name": "value", - "arg": "presentValueT", - "buffer": { "type": "read-only-storage" }, - "elementType": "$inputVec4" - }, - { "name": "output", "arg": "outputT", "buffer": { "type": "storage" }, "elementType": "$inputVec4" }, - { - "name": "params", - "semantic": "kernel.params", - "buffer": { "type": "uniform" }, - "struct": { - "name": "Params", - "fields": [ - { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, - { "name": "kvSeq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" }, - { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, - { "name": "isCausal", "type": "u32", "value": 1 }, - { - "name": "windowSize", - "type": "u32", - "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" - } - ] - } - } - ], - "rotaryFusedFlashAttention": [ - { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputVec4" }, - { "name": "key", "arg": "presentKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputVec4" }, - { - "name": "value", - "arg": "presentValueT", - "buffer": { "type": "read-only-storage" }, - "elementType": "$inputVec4" - }, - { - "name": "cos_cache", - "arg": "cosCacheT", - "buffer": { "type": "read-only-storage" }, - "elementType": "$cosScalar" - }, - { - "name": "sin_cache", - "arg": "sinCacheT", - "buffer": { "type": "read-only-storage" }, - "elementType": "$cosScalar" - }, - { "name": "output", "arg": "outputT", "buffer": { "type": "storage" }, "elementType": "$inputVec4" }, - { - "name": "params", - "semantic": "kernel.params", - "buffer": { "type": "uniform" }, - "struct": { - "name": "Params", - "fields": [ - { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, - { "name": "kvSeq", "type": "u32", "value": "dim(shapes.pastKeyT, 2)" }, - { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, - { "name": "isCausal", "type": "u32", "value": 1 }, - { - "name": "windowSize", - "type": "u32", - "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" - } - ] - } - } - ], - "rotaryFlashAttention": [ - { "name": "query", "semantic": "qPrep", "buffer": { "type": "read-only-storage" }, "elementType": "vec4" }, - { "name": "key", "arg": "presentKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputVec4" }, - { - "name": "value", - "arg": "presentValueT", - "buffer": { "type": "read-only-storage" }, - "elementType": "$inputVec4" - }, - { "name": "output", "arg": "outputT", "buffer": { "type": "storage" }, "elementType": "$inputVec4" }, - { - "name": "params", - "semantic": "kernel.params", - "buffer": { "type": "uniform" }, - "struct": { - "name": "Params", - "fields": [ - { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, - { "name": "kvSeq", "type": "u32", "value": "dim(shapes.pastKeyT, 2)" }, - { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, - { "name": "isCausal", "type": "u32", "value": 1 }, - { - "name": "windowSize", - "type": "u32", - "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" - } - ] - } - } - ], - "qnormRotaryQprep": [ - { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, - { - "name": "cos_cache", - "arg": "cosCacheT", - "buffer": { "type": "read-only-storage" }, - "elementType": "$cosScalar" - }, - { - "name": "sin_cache", - "arg": "sinCacheT", - "buffer": { "type": "read-only-storage" }, - "elementType": "$cosScalar" - }, - { - "name": "q_norm_weight", - "arg": "qNormWeightT", - "buffer": { "type": "read-only-storage" }, - "elementType": "$qNormScalar" - }, - { "name": "qout", "semantic": "qPrep", "buffer": { "type": "storage" }, "elementType": "vec4" }, - { - "name": "params", - "semantic": "kernel.params", - "buffer": { "type": "uniform" }, - "struct": { - "name": "Params", - "fields": [ - { "name": "batch", "type": "u32", "value": "dim(shapes.queryT, 0)" }, - { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, - { "name": "totalSeq", "type": "u32", "value": "dim(shapes.pastKeyT, 2)" } - ] - } - } - ], - "biasFlashAttention": [ - { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputVec4" }, - { "name": "key", "arg": "presentKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputVec4" }, - { - "name": "value", - "arg": "presentValueT", - "buffer": { "type": "read-only-storage" }, - "elementType": "$inputVec4" - }, - { - "name": "attn_mask", - "arg": "attentionBiasT", - "buffer": { "type": "read-only-storage" }, - "elementType": "$inputScalar" - }, - { "name": "output", "arg": "outputT", "buffer": { "type": "storage" }, "elementType": "$inputVec4" }, - { - "name": "params", - "semantic": "kernel.params", - "buffer": { "type": "uniform" }, - "struct": { - "name": "Params", - "fields": [ - { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, - { "name": "kvSeq", "type": "u32", "value": "dim(shapes.pastKeyT, 2)" }, - { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, - { "name": "isCausal", "type": "u32", "value": 1 }, - { - "name": "windowSize", - "type": "u32", - "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" - }, - { - "name": "maskBatchStride", - "type": "u32", - "value": "0 if dim(shapes.attentionBiasT, 0) == 1 else dim(shapes.attentionBiasT, 1) * dim(shapes.attentionBiasT, 2) * dim(shapes.attentionBiasT, 3)" - }, - { - "name": "maskHeadStride", - "type": "u32", - "value": "0 if dim(shapes.attentionBiasT, 1) == 1 else dim(shapes.attentionBiasT, 2) * dim(shapes.attentionBiasT, 3)" - }, - { "name": "maskSeqStride", "type": "u32", "value": "dim(shapes.attentionBiasT, 3)" } - ] - } - } - ], - "qkvSplitAttention": [ - { - "name": "query", - "arg": "queryT", - "semantic": "query", - "buffer": { "type": "read-only-storage" }, - "elementType": "$inputVec4" - }, - { - "name": "key", - "arg": "keyT", - "semantic": "key", - "buffer": { "type": "read-only-storage" }, - "elementType": "$inputVec4" - }, - { - "name": "value", - "arg": "valueT", - "semantic": "value", - "buffer": { "type": "read-only-storage" }, - "elementType": "$inputVec4" - }, - { "name": "partial_out", "semantic": "partialOut", "buffer": { "type": "storage" }, "elementType": "vec4" }, - { - "name": "partial_stats", - "semantic": "partialStats", - "buffer": { "type": "storage" }, - "elementType": "vec2" - }, - { - "name": "params", - "semantic": "kernel.params", - "buffer": { "type": "uniform" }, - "struct": { - "name": "Params", - "fields": [ - { "name": "kvSeq", "type": "u32", "value": "dim(shapes.key, 1)" }, - { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, - { - "name": "windowSize", - "type": "u32", - "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" - } - ] - } - } - ], - "qkvSplitMerge": [ - { - "name": "partial_out", - "semantic": "partialOut", - "buffer": { "type": "read-only-storage" }, - "elementType": "vec4" - }, - { - "name": "partial_stats", - "semantic": "partialStats", - "buffer": { "type": "read-only-storage" }, - "elementType": "vec2" - }, - { - "name": "output", - "arg": "outputT", - "semantic": "output", - "buffer": { "type": "storage" }, - "elementType": "$inputVec4" - }, - { - "name": "params", - "semantic": "kernel.params", - "buffer": { "type": "uniform" }, - "struct": { - "name": "Params", - "fields": [ - { "name": "qSeq", "type": "u32", "value": "dim(shapes.query, 1)" }, - { "name": "kvSeq", "type": "u32", "value": "dim(shapes.key, 1)" }, - { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" } - ] - } - } - ], - "qkvSplitMergeNoParams": [ - { - "name": "partial_out", - "semantic": "partialOut", - "buffer": { "type": "read-only-storage" }, - "elementType": "vec4" - }, - { - "name": "partial_stats", - "semantic": "partialStats", - "buffer": { "type": "read-only-storage" }, - "elementType": "vec2" - }, - { - "name": "output", - "arg": "outputT", - "semantic": "output", - "buffer": { "type": "storage" }, - "elementType": "$inputVec4" - } - ], - "qkvPresent": [ - { - "name": "key", - "arg": "keyT", - "semantic": "key", - "buffer": { "type": "read-only-storage" }, - "elementType": "$inputElement" - }, - { - "name": "value", - "arg": "valueT", - "semantic": "value", - "buffer": { "type": "read-only-storage" }, - "elementType": "$inputElement" - }, - { - "name": "present_key", - "arg": "presentKeyT", - "semantic": "present_key", - "buffer": { "type": "storage" }, - "elementType": "$presentElement" - }, - { - "name": "present_value", - "arg": "presentValueT", - "semantic": "present_value", - "buffer": { "type": "storage" }, - "elementType": "$presentElement" - }, - { - "name": "params", - "semantic": "kernel.params", - "buffer": { "type": "uniform" }, - "struct": { - "name": "Params", - "fields": [ - { "name": "batchSize", "type": "u32", "value": "dim(shapes.query, 0)" }, - { "name": "kvSeq", "type": "u32", "value": "dim(shapes.key, 1)" } - ] - } - } - ], - "qkvAttention": [ - { - "name": "query", - "arg": "queryT", - "semantic": "query", - "buffer": { "type": "read-only-storage" }, - "elementType": "$inputElement" - }, - { - "name": "key", - "arg": "keyT", - "semantic": "key", - "buffer": { "type": "read-only-storage" }, - "elementType": "$inputElement" - }, - { - "name": "value", - "arg": "valueT", - "semantic": "value", - "buffer": { "type": "read-only-storage" }, - "elementType": "$inputElement" - }, - { - "name": "output", - "arg": "outputT", - "semantic": "output", - "buffer": { "type": "storage" }, - "elementType": "$outputElement" - }, - { - "name": "params", - "semantic": "kernel.params", - "buffer": { "type": "uniform" }, - "struct": { - "name": "Params", - "fields": [ - { "name": "qSeq", "type": "u32", "value": "dim(shapes.query, 1)" }, - { "name": "kvSeq", "type": "u32", "value": "dim(shapes.key, 1)" }, - { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, - { - "name": "windowSize", - "type": "u32", - "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" - } - ] - } - } - ], - "qkvAttentionNoWindow": [ - { - "name": "query", - "arg": "queryT", - "semantic": "query", - "buffer": { "type": "read-only-storage" }, - "elementType": "$inputElement" - }, - { - "name": "key", - "arg": "keyT", - "semantic": "key", - "buffer": { "type": "read-only-storage" }, - "elementType": "$inputElement" - }, - { - "name": "value", - "arg": "valueT", - "semantic": "value", - "buffer": { "type": "read-only-storage" }, - "elementType": "$inputElement" - }, - { - "name": "output", - "arg": "outputT", - "semantic": "output", - "buffer": { "type": "storage" }, - "elementType": "$outputElement" - }, - { - "name": "params", - "semantic": "kernel.params", - "buffer": { "type": "uniform" }, - "struct": { - "name": "Params", - "fields": [ - { "name": "qSeq", "type": "u32", "value": "dim(shapes.query, 1)" }, - { "name": "kvSeq", "type": "u32", "value": "dim(shapes.key, 1)" }, - { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" } - ] - } - } - ], - "presentCopy": [ - { "name": "src_k", "arg": "pastKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, - { "name": "src_v", "arg": "pastValueT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, - { "name": "present_key", "arg": "presentKeyT", "buffer": { "type": "storage" }, "elementType": "$inputScalar" }, - { - "name": "present_value", - "arg": "presentValueT", - "buffer": { "type": "storage" }, - "elementType": "$inputScalar" - }, - { - "name": "params", - "semantic": "kernel.params", - "buffer": { "type": "uniform" }, - "struct": { - "name": "Params", - "fields": [ - { - "name": "count", - "type": "u32", - "value": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)" - } - ] - } - } - ], - "presentMergeShare": [ - { "name": "past_k", "arg": "pastKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, - { - "name": "past_v", - "arg": "pastValueT", - "buffer": { "type": "read-only-storage" }, - "elementType": "$inputScalar" - }, - { "name": "new_k", "arg": "keyT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, - { "name": "new_v", "arg": "valueT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, - { "name": "present_key", "arg": "presentKeyT", "buffer": { "type": "storage" }, "elementType": "$inputScalar" }, - { - "name": "present_value", - "arg": "presentValueT", - "buffer": { "type": "storage" }, - "elementType": "$inputScalar" - }, - { "name": "seqlens_k", "arg": "seqlensKT", "buffer": { "type": "read-only-storage" }, "elementType": "i32" }, - { - "name": "params", - "semantic": "kernel.params", - "buffer": { "type": "uniform" }, - "struct": { - "name": "Params", - "fields": [ - { - "name": "count", - "type": "u32", - "value": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)" - }, - { "name": "seq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" }, - { "name": "keySeq", "type": "u32", "value": "dim(shapes.keyT, 1)" } - ] - } - } - ], - "presentMergeShareRetain": [ - { "name": "past_k", "arg": "pastKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, - { - "name": "past_v", - "arg": "pastValueT", - "buffer": { "type": "read-only-storage" }, - "elementType": "$inputScalar" - }, - { "name": "present_key", "arg": "presentKeyT", "buffer": { "type": "storage" }, "elementType": "$inputScalar" }, - { - "name": "present_value", - "arg": "presentValueT", - "buffer": { "type": "storage" }, - "elementType": "$inputScalar" - }, - { "name": "seqlens_k", "arg": "seqlensKT", "buffer": { "type": "read-only-storage" }, "elementType": "i32" }, - { - "name": "params", - "semantic": "kernel.params", - "buffer": { "type": "uniform" }, - "struct": { - "name": "Params", - "fields": [ - { - "name": "count", - "type": "u32", - "value": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)" - }, - { "name": "seq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" }, - { "name": "keySeq", "type": "u32", "value": "dim(shapes.keyT, 1)" } - ] - } - } - ], - "presentMergeShareAppend": [ - { "name": "new_k", "arg": "keyT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, - { "name": "new_v", "arg": "valueT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, - { "name": "present_key", "arg": "presentKeyT", "buffer": { "type": "storage" }, "elementType": "$inputScalar" }, - { - "name": "present_value", - "arg": "presentValueT", - "buffer": { "type": "storage" }, - "elementType": "$inputScalar" - }, - { "name": "seqlens_k", "arg": "seqlensKT", "buffer": { "type": "read-only-storage" }, "elementType": "i32" }, - { - "name": "params", - "semantic": "kernel.params", - "buffer": { "type": "uniform" }, - "struct": { - "name": "Params", - "fields": [ - { - "name": "count", - "type": "u32", - "value": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)" - }, - { "name": "seq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" }, - { "name": "keySeq", "type": "u32", "value": "dim(shapes.keyT, 1)" } - ] - } - } - ], - "presentMergeShareRotaryAppend": [ - { "name": "new_k", "arg": "keyT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, - { "name": "new_v", "arg": "valueT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, - { "name": "cos_cache", "arg": "cosCacheT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, - { "name": "sin_cache", "arg": "sinCacheT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, - { "name": "present_key", "arg": "presentKeyT", "buffer": { "type": "storage" }, "elementType": "$inputScalar" }, - { - "name": "present_value", - "arg": "presentValueT", - "buffer": { "type": "storage" }, - "elementType": "$inputScalar" - }, - { "name": "seqlens_k", "arg": "seqlensKT", "buffer": { "type": "read-only-storage" }, "elementType": "i32" }, - { - "name": "params", - "semantic": "kernel.params", - "buffer": { "type": "uniform" }, - "struct": { - "name": "Params", - "fields": [ - { - "name": "count", - "type": "u32", - "value": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)" - }, - { "name": "seq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" }, - { "name": "keySeq", "type": "u32", "value": "dim(shapes.keyT, 1)" } - ] - } - } - ], - "mergedSeqlensAttention": [ - { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, - { - "name": "kcache", - "arg": "presentKeyT", - "buffer": { "type": "read-only-storage" }, - "elementType": "$inputScalar" - }, - { - "name": "vcache", - "arg": "presentValueT", - "buffer": { "type": "read-only-storage" }, - "elementType": "$inputScalar" - }, - { "name": "output", "arg": "outputT", "buffer": { "type": "storage" }, "elementType": "$inputScalar" }, - { "name": "seqlens_k", "arg": "seqlensKT", "buffer": { "type": "read-only-storage" }, "elementType": "i32" }, - { - "name": "params", - "semantic": "kernel.params", - "buffer": { "type": "uniform" }, - "struct": { - "name": "Params", - "fields": [ - { "name": "batch", "type": "u32", "value": "dim(shapes.queryT, 0)" }, - { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, - { "name": "totalSeq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" }, - { - "name": "windowSize", - "type": "u32", - "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" - }, - { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, - { "name": "softcap", "type": "f32", "value": "attrs.softcap" }, - { "name": "smoothSoftmax", "type": "u32", "value": "1 if attrs.smooth_softmax == 1 else 0" } - ] - } - } - ], - "mergedSeqlensHeadSinkAttention": [ - { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, - { - "name": "kcache", - "arg": "presentKeyT", - "buffer": { "type": "read-only-storage" }, - "elementType": "$inputScalar" - }, - { - "name": "vcache", - "arg": "presentValueT", - "buffer": { "type": "read-only-storage" }, - "elementType": "$inputScalar" - }, - { - "name": "head_sink", - "arg": "headSinkT", - "buffer": { "type": "read-only-storage" }, - "elementType": "$inputScalar" - }, - { "name": "output", "arg": "outputT", "buffer": { "type": "storage" }, "elementType": "$inputScalar" }, - { "name": "seqlens_k", "arg": "seqlensKT", "buffer": { "type": "read-only-storage" }, "elementType": "i32" }, - { - "name": "params", - "semantic": "kernel.params", - "buffer": { "type": "uniform" }, - "struct": { - "name": "Params", - "fields": [ - { "name": "batch", "type": "u32", "value": "dim(shapes.queryT, 0)" }, - { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, - { "name": "totalSeq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" }, - { - "name": "windowSize", - "type": "u32", - "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" - }, - { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, - { "name": "softcap", "type": "f32", "value": "attrs.softcap" } - ] - } - } - ], - "mergedSeqlensRotaryAttention": [ - { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, - { - "name": "kcache", - "arg": "presentKeyT", - "buffer": { "type": "read-only-storage" }, - "elementType": "$inputScalar" - }, - { - "name": "vcache", - "arg": "presentValueT", - "buffer": { "type": "read-only-storage" }, - "elementType": "$inputScalar" - }, - { "name": "cos_cache", "arg": "cosCacheT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, - { "name": "sin_cache", "arg": "sinCacheT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, - { "name": "output", "arg": "outputT", "buffer": { "type": "storage" }, "elementType": "$inputScalar" }, - { "name": "seqlens_k", "arg": "seqlensKT", "buffer": { "type": "read-only-storage" }, "elementType": "i32" }, - { - "name": "params", - "semantic": "kernel.params", - "buffer": { "type": "uniform" }, - "struct": { - "name": "Params", - "fields": [ - { "name": "batch", "type": "u32", "value": "dim(shapes.queryT, 0)" }, - { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, - { "name": "totalSeq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" }, - { - "name": "windowSize", - "type": "u32", - "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" - }, - { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, - { "name": "softcap", "type": "f32", "value": "attrs.softcap" }, - { "name": "smoothSoftmax", "type": "u32", "value": "1 if attrs.smooth_softmax == 1 else 0" } - ] - } - } - ], - "presentMerge": [ - { "name": "past_k", "arg": "pastKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, - { - "name": "past_v", - "arg": "pastValueT", - "buffer": { "type": "read-only-storage" }, - "elementType": "$inputScalar" - }, - { "name": "new_k", "arg": "keyT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, - { "name": "new_v", "arg": "valueT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, - { "name": "present_key", "arg": "presentKeyT", "buffer": { "type": "storage" }, "elementType": "$inputScalar" }, - { - "name": "present_value", - "arg": "presentValueT", - "buffer": { "type": "storage" }, - "elementType": "$inputScalar" - }, - { - "name": "params", - "semantic": "kernel.params", - "buffer": { "type": "uniform" }, - "struct": { - "name": "Params", - "fields": [ - { - "name": "count", - "type": "u32", - "value": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)" - }, - { "name": "seq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" }, - { "name": "pastSeq", "type": "u32", "value": "dim(shapes.pastKeyT, 2)" }, - { "name": "keySeq", "type": "u32", "value": "dim(shapes.keyT, 1)" } - ] - } - } - ], - "qkvTiledAttention": [ - { - "name": "q", - "arg": "queryT", - "semantic": "query", - "buffer": { "type": "read-only-storage" }, - "elementType": "$scalar" - }, - { - "name": "k", - "arg": "keyT", - "semantic": "key", - "buffer": { "type": "read-only-storage" }, - "elementType": "$scalar" - }, - { - "name": "v", - "arg": "valueT", - "semantic": "value", - "buffer": { "type": "read-only-storage" }, - "elementType": "$scalar" - }, - { "name": "y", "arg": "outputT", "semantic": "output", "buffer": { "type": "storage" }, "elementType": "$scalar" }, - { - "name": "params", - "semantic": "kernel.params", - "buffer": { "type": "uniform" }, - "struct": { - "name": "Params", - "fields": [ - { "name": "count", "type": "u32", "value": "numel(shapes.outputT)" }, - { "name": "qHeads", "type": "u32", "value": "attrs.num_heads" }, - { "name": "kvHeads", "type": "u32", "value": "attrs.kv_num_heads" }, - { "name": "qSeq", "type": "u32", "value": "dim(shapes.query, 1)" }, - { "name": "kvSeq", "type": "u32", "value": "dim(shapes.key, 1)" }, - { "name": "headSize", "type": "u32", "value": "dim(shapes.query, 2) / attrs.num_heads" }, - { "name": "vHeadSize", "type": "u32", "value": "dim(shapes.value, 2) / attrs.kv_num_heads" }, - { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, - { "name": "softcap", "type": "f32", "value": "0" }, - { "name": "isCausal", "type": "u32", "value": "0" }, - { "name": "qHidden", "type": "u32", "value": "dim(shapes.query, 2)" }, - { "name": "kvHidden", "type": "u32", "value": "dim(shapes.key, 2)" }, - { "name": "vHidden", "type": "u32", "value": "dim(shapes.value, 2)" } - ] - } - } - ], - "cachedAttention": [ - { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, - { "name": "kcache", "arg": "pastKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, - { - "name": "vcache", - "arg": "pastValueT", - "buffer": { "type": "read-only-storage" }, - "elementType": "$inputScalar" - }, - { "name": "output", "arg": "outputT", "buffer": { "type": "storage" }, "elementType": "$inputScalar" }, - { "name": "seqlens_k", "arg": "seqlensKT", "buffer": { "type": "read-only-storage" }, "elementType": "i32" }, - { - "name": "params", - "semantic": "kernel.params", - "buffer": { "type": "uniform" }, - "struct": { - "name": "Params", - "fields": [ - { "name": "batch", "type": "u32", "value": "dim(shapes.queryT, 0)" }, - { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, - { "name": "totalSeq", "type": "u32", "value": "dim(shapes.pastKeyT, 2)" }, - { - "name": "windowSize", - "type": "u32", - "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" - }, - { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, - { "name": "softcap", "type": "f32", "value": "attrs.softcap" }, - { "name": "smoothSoftmax", "type": "u32", "value": "1 if attrs.smooth_softmax == 1 else 0" } - ] - } - } - ], - "cachedRotaryAttention": [ - { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, - { "name": "kcache", "arg": "pastKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, - { - "name": "vcache", - "arg": "pastValueT", - "buffer": { "type": "read-only-storage" }, - "elementType": "$inputScalar" - }, - { - "name": "cos_cache", - "arg": "cosCacheT", - "buffer": { "type": "read-only-storage" }, - "elementType": "$cosScalar" - }, - { - "name": "sin_cache", - "arg": "sinCacheT", - "buffer": { "type": "read-only-storage" }, - "elementType": "$cosScalar" - }, - { "name": "output", "arg": "outputT", "buffer": { "type": "storage" }, "elementType": "$inputScalar" }, - { "name": "seqlens_k", "arg": "seqlensKT", "buffer": { "type": "read-only-storage" }, "elementType": "i32" }, - { - "name": "params", - "semantic": "kernel.params", - "buffer": { "type": "uniform" }, - "struct": { - "name": "Params", - "fields": [ - { "name": "batch", "type": "u32", "value": "dim(shapes.queryT, 0)" }, - { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, - { "name": "totalSeq", "type": "u32", "value": "dim(shapes.pastKeyT, 2)" }, - { - "name": "windowSize", - "type": "u32", - "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" - }, - { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, - { "name": "softcap", "type": "f32", "value": "attrs.softcap" }, - { "name": "smoothSoftmax", "type": "u32", "value": "1 if attrs.smooth_softmax == 1 else 0" } - ] - } - } - ], - "cachedQnormRotaryAttention": [ - { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, - { "name": "kcache", "arg": "pastKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, - { - "name": "vcache", - "arg": "pastValueT", - "buffer": { "type": "read-only-storage" }, - "elementType": "$inputScalar" - }, - { - "name": "cos_cache", - "arg": "cosCacheT", - "buffer": { "type": "read-only-storage" }, - "elementType": "$cosScalar" - }, - { - "name": "sin_cache", - "arg": "sinCacheT", - "buffer": { "type": "read-only-storage" }, - "elementType": "$cosScalar" - }, - { - "name": "q_norm_weight", - "arg": "qNormWeightT", - "buffer": { "type": "read-only-storage" }, - "elementType": "$qNormScalar" - }, - { "name": "output", "arg": "outputT", "buffer": { "type": "storage" }, "elementType": "$inputScalar" }, - { "name": "seqlens_k", "arg": "seqlensKT", "buffer": { "type": "read-only-storage" }, "elementType": "i32" }, - { - "name": "params", - "semantic": "kernel.params", - "buffer": { "type": "uniform" }, - "struct": { - "name": "Params", - "fields": [ - { "name": "batch", "type": "u32", "value": "dim(shapes.queryT, 0)" }, - { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, - { "name": "totalSeq", "type": "u32", "value": "dim(shapes.pastKeyT, 2)" }, - { - "name": "windowSize", - "type": "u32", - "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" - }, - { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, - { "name": "softcap", "type": "f32", "value": "attrs.softcap" }, - { "name": "smoothSoftmax", "type": "u32", "value": "1 if attrs.smooth_softmax == 1 else 0" } - ] - } - } - ], - "newQnormRotaryPresent": [ - { "name": "src_k", "arg": "keyT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, - { "name": "src_v", "arg": "valueT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, - { "name": "cos_cache", "arg": "cosCacheT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, - { "name": "sin_cache", "arg": "sinCacheT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, - { - "name": "k_norm_weight", - "arg": "kNormWeightT", - "buffer": { "type": "read-only-storage" }, - "elementType": "f32" - }, - { "name": "present_key", "arg": "presentKeyT", "buffer": { "type": "storage" }, "elementType": "f32" }, - { "name": "present_value", "arg": "presentValueT", "buffer": { "type": "storage" }, "elementType": "f32" }, - { - "name": "params", - "semantic": "kernel.params", - "buffer": { "type": "uniform" }, - "struct": { - "name": "Params", - "fields": [ - { - "name": "count", - "type": "u32", - "value": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)" - }, - { "name": "seq", "type": "u32", "value": "dim(shapes.keyT, 1)" }, - { "name": "pastSeq", "type": "u32", "value": 0 } - ] - } - } - ], - "newQnormRotaryAttention": [ - { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, - { "name": "kcache", "arg": "presentKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, - { "name": "vcache", "arg": "presentValueT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, - { "name": "cos_cache", "arg": "cosCacheT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, - { "name": "sin_cache", "arg": "sinCacheT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, - { - "name": "q_norm_weight", - "arg": "qNormWeightT", - "buffer": { "type": "read-only-storage" }, - "elementType": "f32" - }, - { "name": "output", "arg": "outputT", "buffer": { "type": "storage" }, "elementType": "f32" }, - { - "name": "params", - "semantic": "kernel.params", - "buffer": { "type": "uniform" }, - "struct": { - "name": "Params", - "fields": [ - { "name": "batch", "type": "u32", "value": "dim(shapes.queryT, 0)" }, - { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, - { "name": "totalSeq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" }, - { - "name": "windowSize", - "type": "u32", - "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" - }, - { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, - { "name": "softcap", "type": "f32", "value": "attrs.softcap" }, - { "name": "smoothSoftmax", "type": "u32", "value": "1 if attrs.smooth_softmax == 1 else 0" } - ] - } - } - ], - "mergedAttention": [ - { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" }, - { - "name": "kcache", - "arg": "presentKeyT", - "buffer": { "type": "read-only-storage" }, - "elementType": "$inputScalar" - }, - { - "name": "vcache", - "arg": "presentValueT", - "buffer": { "type": "read-only-storage" }, - "elementType": "$inputScalar" - }, - { "name": "output", "arg": "outputT", "buffer": { "type": "storage" }, "elementType": "$inputScalar" }, - { - "name": "params", - "semantic": "kernel.params", - "buffer": { "type": "uniform" }, - "struct": { - "name": "Params", - "fields": [ - { "name": "batch", "type": "u32", "value": "dim(shapes.queryT, 0)" }, - { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, - { "name": "totalSeq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" }, - { - "name": "windowSize", - "type": "u32", - "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" - }, - { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, - { "name": "softcap", "type": "f32", "value": "attrs.softcap" }, - { "name": "smoothSoftmax", "type": "u32", "value": "1 if attrs.smooth_softmax == 1 else 0" } - ] - } - } - ], - "quantPresent": [ - { "name": "src_k", "arg": "keyT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, - { "name": "src_v", "arg": "valueT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, - { "name": "k_scale", "arg": "kScaleT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, - { "name": "v_scale", "arg": "vScaleT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, - { - "name": "present_key", - "arg": "presentKeyT", - "buffer": { "type": "storage" }, - "elementType": "$quantCacheElement" - }, - { - "name": "present_value", - "arg": "presentValueT", - "buffer": { "type": "storage" }, - "elementType": "$quantCacheElement" - }, - { - "name": "params", - "semantic": "kernel.params", - "buffer": { "type": "uniform" }, - "struct": { - "name": "Params", - "fields": [ - { - "name": "count", - "type": "u32", - "value": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)" - }, - { "name": "seq", "type": "u32", "value": "dim(shapes.keyT, 1)" }, - { "name": "perChannel", "type": "u32", "value": "1 if dim(shapes.kScaleT, 0) > 1 else 0" } - ] - } - } - ], - "quantInt8DecodePresent": [ - { "name": "src_k", "arg": "keyT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, - { "name": "src_v", "arg": "valueT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, - { "name": "k_scale", "arg": "kScaleT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, - { "name": "v_scale", "arg": "vScaleT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, - { "name": "seqlens_k", "arg": "seqlensKT", "buffer": { "type": "read-only-storage" }, "elementType": "i32" }, - { "name": "present_key", "arg": "presentKeyT", "buffer": { "type": "storage" }, "elementType": "i32" }, - { "name": "present_value", "arg": "presentValueT", "buffer": { "type": "storage" }, "elementType": "i32" }, - { - "name": "params", - "semantic": "kernel.params", - "buffer": { "type": "uniform" }, - "struct": { - "name": "Params", - "fields": [ - { - "name": "count", - "type": "u32", - "value": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)" - }, - { "name": "seq", "type": "u32", "value": "dim(shapes.keyT, 1)" }, - { "name": "totalSeq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" }, - { "name": "perChannel", "type": "u32", "value": "1 if dim(shapes.kScaleT, 0) > 1 else 0" } - ] - } - } - ], - "quantAttention": [ - { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, - { - "name": "kcache", - "arg": "presentKeyT", - "buffer": { "type": "read-only-storage" }, - "elementType": "$quantCacheElement" - }, - { - "name": "vcache", - "arg": "presentValueT", - "buffer": { "type": "read-only-storage" }, - "elementType": "$quantCacheElement" - }, - { "name": "k_scale", "arg": "kScaleT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, - { "name": "v_scale", "arg": "vScaleT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, - { "name": "output", "arg": "outputT", "buffer": { "type": "storage" }, "elementType": "f32" }, - { - "name": "params", - "semantic": "kernel.params", - "buffer": { "type": "uniform" }, - "struct": { - "name": "Params", - "fields": [ - { "name": "batch", "type": "u32", "value": "dim(shapes.queryT, 0)" }, - { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, - { "name": "totalSeq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" }, - { - "name": "windowSize", - "type": "u32", - "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" - }, - { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, - { "name": "softcap", "type": "f32", "value": "attrs.softcap" }, - { "name": "smoothSoftmax", "type": "u32", "value": "1 if attrs.smooth_softmax == 1 else 0" }, - { "name": "perChannel", "type": "u32", "value": "1 if dim(shapes.kScaleT, 0) > 1 else 0" } - ] - } - } - ], - "quantInt8SplitAttention": [ - { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "vec4" }, - { "name": "key", "arg": "presentKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "vec4" }, - { "name": "value", "arg": "presentValueT", "buffer": { "type": "read-only-storage" }, "elementType": "vec4" }, - { "name": "k_scale", "arg": "kScaleT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, - { "name": "v_scale", "arg": "vScaleT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, - { "name": "seqlens_k", "arg": "seqlensKT", "buffer": { "type": "read-only-storage" }, "elementType": "i32" }, - { "name": "partial_out", "semantic": "partialOut", "buffer": { "type": "storage" }, "elementType": "vec4" }, - { - "name": "partial_stats", - "semantic": "partialStats", - "buffer": { "type": "storage" }, - "elementType": "vec2" - }, - { - "name": "params", - "semantic": "kernel.params", - "buffer": { "type": "uniform" }, - "struct": { - "name": "Params", - "fields": [ - { "name": "kvSeq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" }, - { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, - { - "name": "windowSize", - "type": "u32", - "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" - }, - { "name": "perChannel", "type": "u32", "value": "1 if dim(shapes.kScaleT, 0) > 1 else 0" } - ] - } - } - ], - "cachedBiasAttention": [ - { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, - { "name": "kcache", "arg": "pastKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, - { "name": "vcache", "arg": "pastValueT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, - { - "name": "attn_bias", - "arg": "attentionBiasT", - "buffer": { "type": "read-only-storage" }, - "elementType": "$inputScalar" - }, - { "name": "output", "arg": "outputT", "buffer": { "type": "storage" }, "elementType": "f32" }, - { "name": "seqlens_k", "arg": "seqlensKT", "buffer": { "type": "read-only-storage" }, "elementType": "i32" }, - { - "name": "params", - "semantic": "kernel.params", - "buffer": { "type": "uniform" }, - "struct": { - "name": "Params", - "fields": [ - { "name": "batch", "type": "u32", "value": "dim(shapes.queryT, 0)" }, - { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, - { "name": "totalSeq", "type": "u32", "value": "dim(shapes.pastKeyT, 2)" }, - { - "name": "windowSize", - "type": "u32", - "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" - }, - { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, - { "name": "softcap", "type": "f32", "value": "attrs.softcap" }, - { "name": "smoothSoftmax", "type": "u32", "value": "1 if attrs.smooth_softmax == 1 else 0" }, - { "name": "biasBatch", "type": "u32", "value": "dim(shapes.attentionBiasT, 0)" }, - { "name": "biasHeads", "type": "u32", "value": "dim(shapes.attentionBiasT, 1)" } - ] - } - } - ], - "cachedHeadSinkAttention": [ - { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, - { "name": "kcache", "arg": "pastKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, - { "name": "vcache", "arg": "pastValueT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, - { - "name": "head_sink", - "arg": "headSinkT", - "buffer": { "type": "read-only-storage" }, - "elementType": "$inputScalar" - }, - { "name": "output", "arg": "outputT", "buffer": { "type": "storage" }, "elementType": "f32" }, - { "name": "seqlens_k", "arg": "seqlensKT", "buffer": { "type": "read-only-storage" }, "elementType": "i32" }, - { - "name": "params", - "semantic": "kernel.params", - "buffer": { "type": "uniform" }, - "struct": { - "name": "Params", - "fields": [ - { "name": "batch", "type": "u32", "value": "dim(shapes.queryT, 0)" }, - { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, - { "name": "totalSeq", "type": "u32", "value": "dim(shapes.pastKeyT, 2)" }, - { - "name": "windowSize", - "type": "u32", - "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" - }, - { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, - { "name": "softcap", "type": "f32", "value": "attrs.softcap" } - ] - } - } - ], - "cachedBiasHeadSinkAttention": [ - { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, - { "name": "kcache", "arg": "pastKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, - { "name": "vcache", "arg": "pastValueT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, - { - "name": "attn_bias", - "arg": "attentionBiasT", - "buffer": { "type": "read-only-storage" }, - "elementType": "$inputScalar" - }, - { - "name": "head_sink", - "arg": "headSinkT", - "buffer": { "type": "read-only-storage" }, - "elementType": "$inputScalar" - }, - { "name": "output", "arg": "outputT", "buffer": { "type": "storage" }, "elementType": "f32" }, - { "name": "seqlens_k", "arg": "seqlensKT", "buffer": { "type": "read-only-storage" }, "elementType": "i32" }, - { - "name": "params", - "semantic": "kernel.params", - "buffer": { "type": "uniform" }, - "struct": { - "name": "Params", - "fields": [ - { "name": "batch", "type": "u32", "value": "dim(shapes.queryT, 0)" }, - { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, - { "name": "totalSeq", "type": "u32", "value": "dim(shapes.pastKeyT, 2)" }, - { - "name": "windowSize", - "type": "u32", - "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" - }, - { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, - { "name": "softcap", "type": "f32", "value": "attrs.softcap" }, - { "name": "biasBatch", "type": "u32", "value": "dim(shapes.attentionBiasT, 0)" }, - { "name": "biasHeads", "type": "u32", "value": "dim(shapes.attentionBiasT, 1)" } - ] - } - } - ], - "flashPrefillQuant": [ - { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputVec4" }, - { - "name": "key", - "arg": "presentKeyT", - "buffer": { "type": "read-only-storage" }, - "elementType": "$quantCacheReadElement" - }, - { - "name": "value", - "arg": "presentValueT", - "buffer": { "type": "read-only-storage" }, - "elementType": "$quantCacheReadElement" - }, - { "name": "k_scale", "arg": "kScaleT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, - { "name": "v_scale", "arg": "vScaleT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, - { "name": "output", "arg": "outputT", "buffer": { "type": "storage" }, "elementType": "$inputVec4" }, - { - "name": "params", - "semantic": "kernel.params", - "buffer": { "type": "uniform" }, - "struct": { - "name": "Params", - "fields": [ - { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, - { "name": "kvSeq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" }, - { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, - { "name": "isCausal", "type": "u32", "value": 1 }, - { - "name": "windowSize", - "type": "u32", - "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" - }, - { "name": "perChannel", "type": "u32", "value": "1 if dim(shapes.kScaleT, 0) > 1 else 0" } - ] - } - } - ], - "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" }, - { - "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" } - ] + "bindings": { + "query": { "arg": "queryT", "buffer": "read-only-storage", "elementType": "f32" }, + "key": { "arg": "keyT", "buffer": "read-only-storage", "elementType": "f32" }, + "scores": { "scratch": "materializedScores", "buffer": "storage", "elementType": "f32" }, + "scorePartials": { "scratch": "materializedScorePartials", "buffer": "storage", "elementType": "f32" }, + "scorePartials_2": { + "scratch": "materializedScorePartials", + "name": "scorePartials", + "buffer": "read-only-storage", + "elementType": "f32" + }, + "rowStats": { "scratch": "materializedRowStats", "buffer": "storage", "elementType": "f32" }, + "scores_2": { + "scratch": "materializedScores", + "name": "scores", + "buffer": "read-only-storage", + "elementType": "f32" + }, + "value": { "arg": "valueT", "buffer": "read-only-storage", "elementType": "f32" }, + "rowStats_2": { + "scratch": "materializedRowStats", + "name": "rowStats", + "buffer": "read-only-storage", + "elementType": "f32" + }, + "output": { "arg": "outputT", "buffer": "storage", "elementType": "f32" }, + "key_2": { "arg": "keyT", "name": "key", "buffer": "read-only-storage", "elementType": "$inputElement" }, + "value_2": { "arg": "valueT", "name": "value", "buffer": "read-only-storage", "elementType": "$inputElement" }, + "present_key": { "arg": "presentKeyT", "buffer": "storage", "elementType": "$presentElement" }, + "present_value": { "arg": "presentValueT", "buffer": "storage", "elementType": "$presentElement" }, + "params": { + "buffer": "uniform", + "struct": [ + { "name": "batchSize", "type": "u32", "value": "dim(shapes.queryT, 0)" }, + { "name": "kvSeq", "type": "u32", "value": "dim(shapes.keyT, 1)" } + ] + }, + "src_k": { "arg": "pastKeyT", "buffer": "read-only-storage", "elementType": "$inputScalar" }, + "src_v": { "arg": "pastValueT", "buffer": "read-only-storage", "elementType": "$inputScalar" }, + "present_key_2": { "arg": "presentKeyT", "name": "present_key", "buffer": "storage", "elementType": "$inputScalar" }, + "present_value_2": { + "arg": "presentValueT", + "name": "present_value", + "buffer": "storage", + "elementType": "$inputScalar" + }, + "params_5": { + "name": "params", + "buffer": "uniform", + "struct": [ + { + "name": "count", + "type": "u32", + "value": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)" } - } - ], - "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)" } - ] + ] + }, + "key_3": { "arg": "presentKeyT", "name": "key", "buffer": "read-only-storage", "elementType": "f32" }, + "params_6": { + "name": "params", + "buffer": "uniform", + "struct": [ + { "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" } + ] + }, + "params_7": { + "name": "params", + "buffer": "uniform", + "struct": [ + { "name": "rows", "type": "u32", "value": "dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)" } + ] + }, + "value_3": { "arg": "presentValueT", "name": "value", "buffer": "read-only-storage", "elementType": "f32" }, + "params_8": { + "name": "params", + "buffer": "uniform", + "struct": [ + { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, + { "name": "kvSeq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" } + ] + }, + "key_4": { "arg": "presentKeyT", "name": "key", "buffer": "read-only-storage", "elementType": "f16" }, + "value_4": { "arg": "presentValueT", "name": "value", "buffer": "read-only-storage", "elementType": "f16" }, + "output_2": { "arg": "outputT", "name": "output", "buffer": "storage", "elementType": "f16" }, + "query_3": { "arg": "queryT", "name": "query", "buffer": "read-only-storage", "elementType": "$inputScalar" }, + "cos_cache": { "arg": "cosCacheT", "buffer": "read-only-storage", "elementType": "$cosScalar" }, + "sin_cache": { "arg": "sinCacheT", "buffer": "read-only-storage", "elementType": "$cosScalar" }, + "qout": { "scratch": "qPrep", "buffer": "storage", "elementType": "vec4" }, + "params_10": { + "name": "params", + "buffer": "uniform", + "struct": [ + { "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)" } + ] + }, + "past_k": { "arg": "pastKeyT", "buffer": "read-only-storage", "elementType": "$inputScalar" }, + "past_v": { "arg": "pastValueT", "buffer": "read-only-storage", "elementType": "$inputScalar" }, + "new_k": { "arg": "keyT", "buffer": "read-only-storage", "elementType": "$inputScalar" }, + "new_v": { "arg": "valueT", "buffer": "read-only-storage", "elementType": "$inputScalar" }, + "params_11": { + "name": "params", + "buffer": "uniform", + "struct": [ + { + "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)" } + ] + }, + "seqlens_k": { "arg": "seqlensKT", "buffer": "read-only-storage", "elementType": "i32" }, + "params_12": { + "name": "params", + "buffer": "uniform", + "struct": [ + { + "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)" } + ] + }, + "params_15": { + "name": "params", + "buffer": "uniform", + "struct": [ + { "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)" } + ] + }, + "kcache": { "arg": "presentKeyT", "buffer": "read-only-storage", "elementType": "$inputScalar" }, + "vcache": { "arg": "presentValueT", "buffer": "read-only-storage", "elementType": "$inputScalar" }, + "output_3": { "arg": "outputT", "name": "output", "buffer": "storage", "elementType": "$inputScalar" }, + "params_16": { + "name": "params", + "buffer": "uniform", + "struct": [ + { "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" } + ] + }, + "head_sink": { "arg": "headSinkT", "buffer": "read-only-storage", "elementType": "$inputScalar" }, + "cos_cache_2": { "arg": "cosCacheT", "name": "cos_cache", "buffer": "read-only-storage", "elementType": "f32" }, + "sin_cache_2": { "arg": "sinCacheT", "name": "sin_cache", "buffer": "read-only-storage", "elementType": "f32" }, + "query_6": { "arg": "queryT", "name": "query", "buffer": "read-only-storage", "elementType": "$inputElement" }, + "output_4": { "arg": "outputT", "name": "output", "buffer": "storage", "elementType": "$outputElement" }, + "params_19": { + "name": "params", + "buffer": "uniform", + "struct": [ + { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, + { "name": "kvSeq", "type": "u32", "value": "dim(shapes.keyT, 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" } - } - ], - "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" } - ] + ] + }, + "src_k_2": { "arg": "keyT", "name": "src_k", "buffer": "read-only-storage", "elementType": "f32" }, + "src_v_2": { "arg": "valueT", "name": "src_v", "buffer": "read-only-storage", "elementType": "f32" }, + "k_scale": { "arg": "kScaleT", "buffer": "read-only-storage", "elementType": "f32" }, + "v_scale": { "arg": "vScaleT", "buffer": "read-only-storage", "elementType": "f32" }, + "present_key_3": { + "arg": "presentKeyT", + "name": "present_key", + "buffer": "storage", + "elementType": "$quantCacheElement" + }, + "present_value_3": { + "arg": "presentValueT", + "name": "present_value", + "buffer": "storage", + "elementType": "$quantCacheElement" + }, + "params_20": { + "name": "params", + "buffer": "uniform", + "struct": [ + { "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" } + ] + }, + "partial_out": { "scratch": "partialOut", "buffer": "storage", "elementType": "vec4" }, + "partial_stats": { "scratch": "partialStats", "buffer": "storage", "elementType": "vec2" }, + "partial_out_2": { + "scratch": "partialOut", + "name": "partial_out", + "buffer": "read-only-storage", + "elementType": "vec4" + }, + "partial_stats_2": { + "scratch": "partialStats", + "name": "partial_stats", + "buffer": "read-only-storage", + "elementType": "vec2" + }, + "output_5": { "arg": "outputT", "name": "output", "buffer": "storage", "elementType": "$inputVec4" }, + "query_8": { "arg": "queryT", "name": "query", "buffer": "read-only-storage", "elementType": "$inputVec4" }, + "key_7": { "arg": "presentKeyT", "name": "key", "buffer": "read-only-storage", "elementType": "$inputVec4" }, + "value_7": { "arg": "presentValueT", "name": "value", "buffer": "read-only-storage", "elementType": "$inputVec4" }, + "attn_mask": { "arg": "attentionBiasT", "buffer": "read-only-storage", "elementType": "$inputScalar" }, + "params_25": { + "name": "params", + "buffer": "uniform", + "struct": [ + { "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)" } + ] + }, + "q_norm_weight": { "arg": "qNormWeightT", "buffer": "read-only-storage", "elementType": "$qNormScalar" }, + "params_26": { + "name": "params", + "buffer": "uniform", + "struct": [ + { "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" } - } - ], - "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" }, - { - "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)" } - ] + ] + }, + "params_27": { + "name": "params", + "buffer": "uniform", + "struct": [ + { "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" } - } - ], - "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)" } - ] + ] + }, + "params_28": { + "name": "params", + "buffer": "uniform", + "struct": [ + { "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" } - } - ] + ] + }, + "params_29": { + "name": "params", + "buffer": "uniform", + "struct": [ + { "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" } + ] + }, + "kcache_3": { "arg": "pastKeyT", "name": "kcache", "buffer": "read-only-storage", "elementType": "$inputScalar" }, + "vcache_3": { "arg": "pastValueT", "name": "vcache", "buffer": "read-only-storage", "elementType": "$inputScalar" }, + "params_31": { + "name": "params", + "buffer": "uniform", + "struct": [ + { "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" } + ] + }, + "k_norm_weight": { "arg": "kNormWeightT", "buffer": "read-only-storage", "elementType": "f32" }, + "present_key_5": { "arg": "presentKeyT", "name": "present_key", "buffer": "storage", "elementType": "f32" }, + "present_value_5": { "arg": "presentValueT", "name": "present_value", "buffer": "storage", "elementType": "f32" }, + "params_32": { + "name": "params", + "buffer": "uniform", + "struct": [ + { "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 } + ] + }, + "kcache_4": { "arg": "presentKeyT", "name": "kcache", "buffer": "read-only-storage", "elementType": "f32" }, + "vcache_4": { "arg": "presentValueT", "name": "vcache", "buffer": "read-only-storage", "elementType": "f32" }, + "q_norm_weight_2": { + "arg": "qNormWeightT", + "name": "q_norm_weight", + "buffer": "read-only-storage", + "elementType": "f32" + }, + "kcache_5": { "arg": "pastKeyT", "name": "kcache", "buffer": "read-only-storage", "elementType": "f32" }, + "vcache_5": { "arg": "pastValueT", "name": "vcache", "buffer": "read-only-storage", "elementType": "f32" }, + "attn_bias": { "arg": "attentionBiasT", "buffer": "read-only-storage", "elementType": "$inputScalar" }, + "params_33": { + "name": "params", + "buffer": "uniform", + "struct": [ + { "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)" } + ] + }, + "params_34": { + "name": "params", + "buffer": "uniform", + "struct": [ + { "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" } + ] + }, + "params_35": { + "name": "params", + "buffer": "uniform", + "struct": [ + { "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)" } + ] + }, + "key_8": { + "arg": "presentKeyT", + "name": "key", + "buffer": "read-only-storage", + "elementType": "$quantCacheReadElement" + }, + "value_8": { + "arg": "presentValueT", + "name": "value", + "buffer": "read-only-storage", + "elementType": "$quantCacheReadElement" + }, + "params_36": { + "name": "params", + "buffer": "uniform", + "struct": [ + { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }, + { "name": "kvSeq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" }, + { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, + { "name": "isCausal", "type": "u32", "value": 1 }, + { + "name": "windowSize", + "type": "u32", + "value": "attrs.local_window_size if attrs.local_window_size > 0 else 0" + }, + { "name": "perChannel", "type": "u32", "value": "1 if dim(shapes.kScaleT, 0) > 1 else 0" } + ] + } }, "variants": [ { "id": "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, + "when": ["gqaMatCoreOk"], "requires": { "features": ["subgroups", "chromium-experimental-subgroup-matrix"], "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] }, - "when": ["gqaMatCoreOk"], - "constants": { + "derive": { "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)", + "headDim": "dim(shapes.queryT, 2) / attrs.num_heads", + "qHidden": "dim(shapes.queryT, 2)", + "kvHidden": "dim(shapes.keyT, 2)", "hasBias": false, "useSubgroups": true, "attentionScaleExpression": "\"select(inverseSqrt(f32(HEAD_DIM)), params.scale, params.scale != 0.0)\"", @@ -2386,20 +684,19 @@ "materializedSgmatRuntimeDirectStore": false, "materializedRowStatsWg": "gqaMatRowStatsWg", "statSlots": "gqaMatStatSlots", - "statQuerySeq": "dim(shapes.query, 1)", + "statQuerySeq": "dim(shapes.queryT, 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" + "headDimV4": "(dim(shapes.queryT, 2) / attrs.num_heads) / 4", + "kvHiddenV4": "dim(shapes.keyT, 2) / 4" }, "intermediates": [ { "id": "materializedScores", "dtype": "float32", - "shape": "[dim(shapes.query, 0) * attrs.num_heads * dim(shapes.query, 1) * dim(shapes.key, 1)]" + "shape": "[dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1) * dim(shapes.keyT, 1)]" }, { "id": "materializedRowStats", "dtype": "float32", "shape": "[gqaMatRowStatsElements]" }, { "id": "materializedScorePartials", "dtype": "float32", "shape": "[gqaMatScorePartialElements]" } @@ -2408,60 +705,99 @@ { "id": "scores", "name": "GroupQueryAttention.MaterializedScoresSgmat", - "source": { - "shader": "attn-materialized-sgmat-f32.wgsl.jinja", - "inputs": { "phase": "\"score\"", "emitRowStats": true, "layout": "\"bsh\"" } - }, - "bindings": "gqaMatScoreStats", + "shader": "attn-materialized-sgmat-f32.wgsl.jinja", + "derive": { "phase": "\"score\"", "emitRowStats": true, "layout": "\"bsh\"" }, + "bindings": [ + "query", + "key", + "scores", + "scorePartials", + { + "name": "params", + "struct": [ + { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, + { "name": "kvSeq", "type": "u32", "value": "dim(shapes.keyT, 1)" }, + { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" } + ] + } + ], "dispatch": { - "x": "ceilDiv(dim(shapes.key, 1), gqaMatKeyTile)", - "y": "ceilDiv(dim(shapes.query, 1), gqaMatQueryTile)", - "z": "dim(shapes.query, 0) * attrs.num_heads" - } + "x": "ceilDiv(dim(shapes.keyT, 1), gqaMatKeyTile)", + "y": "ceilDiv(dim(shapes.queryT, 1), gqaMatQueryTile)", + "z": "dim(shapes.queryT, 0) * attrs.num_heads" + }, + "subgroupCollectivesWidth": 32 }, { "id": "rowstats", "name": "GroupQueryAttention.MaterializedRowStatsCombine", - "source": { "shader": "attn-materialized-rowstats-combine-f32.wgsl.jinja" }, - "bindings": "gqaMatRowStatsCombine", + "shader": "attn-materialized-rowstats-combine-f32.wgsl.jinja", + "bindings": [ + "scorePartials_2", + "rowStats", + { + "name": "params", + "struct": [ + { + "name": "rows", + "type": "u32", + "value": "dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)" + } + ] + } + ], "dispatch": { - "gridStride": "dim(shapes.query, 0) * attrs.num_heads * dim(shapes.query, 1)", - "workgroupSize": "gqaMatRowStatsWg" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)), (gqaMatRowStatsWg)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", + "y": 1, + "z": 1 } }, { "id": "apply", "name": "GroupQueryAttention.MaterializedApplySgmat", - "source": { - "shader": "attn-materialized-sgmat-f32.wgsl.jinja", - "inputs": { "phase": "\"apply\"", "fusedSoftmax": true, "layout": "\"bsh\"" } - }, - "bindings": "gqaMatApplyFused", + "shader": "attn-materialized-sgmat-f32.wgsl.jinja", + "derive": { "phase": "\"apply\"", "fusedSoftmax": true, "layout": "\"bsh\"" }, + "bindings": [ + "scores_2", + "value", + "rowStats_2", + "output", + { + "name": "params", + "struct": [ + { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, + { "name": "kvSeq", "type": "u32", "value": "dim(shapes.keyT, 1)" } + ] + } + ], "dispatch": { "x": "ceilDiv(headDim, gqaMatKeyTile)", - "y": "ceilDiv(dim(shapes.query, 1), gqaMatQueryTile)", - "z": "dim(shapes.query, 0) * attrs.num_heads" + "y": "ceilDiv(dim(shapes.queryT, 1), gqaMatQueryTile)", + "z": "dim(shapes.queryT, 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" } + "bindings": ["key_2", "value_2", "present_key", "present_value", "params"], + "dispatch": { + "x": "min(ceilDiv((numel(shapes.presentKeyT) / 4), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((numel(shapes.presentKeyT) / 4), (copyWorkgroupSize)), 65535)", + "z": 1 + } } ] }, { "id": "past_kv_materialized_sgmat_f32", + "priority": 35, + "when": ["gqaMatPastCoreOk"], "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": { + "derive": { "qNumHeads": "attrs.num_heads", "kvNumHeads": "attrs.kv_num_heads", "headDim": "gqaHeadDim", @@ -2482,8 +818,7 @@ "packed": "gqaHeadDim", "kvHeads": "attrs.kv_num_heads", "inputScalar": "gqaScalar", - "usesF16": false, - "copyWorkgroupSize": "copyWorkgroupSize" + "usesF16": false }, "intermediates": [ { @@ -2499,57 +834,56 @@ "id": "present", "name": "GroupQueryAttention.Present", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentCopy", + "bindings": ["src_k", "src_v", "present_key_2", "present_value_2", "params_5"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "z": 1 }, "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 - } + "shader": "attn-materialized-sgmat-f32.wgsl.jinja", + "derive": { + "phase": "\"score\"", + "emitRowStats": true, + "layout": "\"bsh\"", + "kvLayout": "\"bhsd\"", + "causalRightAlign": true }, - "bindings": "gqaMatPastScoreStats", + "bindings": ["query", "key_3", "scores", "scorePartials", "params_6"], "dispatch": { "x": "ceilDiv(dim(shapes.presentKeyT, 2), gqaMatKeyTile)", "y": "ceilDiv(dim(shapes.queryT, 1), gqaMatQueryTile)", "z": "dim(shapes.queryT, 0) * attrs.num_heads" - } + }, + "subgroupCollectivesWidth": 32 }, { "id": "rowstats", "name": "GroupQueryAttention.PastMaterializedRowStats", - "source": { "shader": "attn-materialized-rowstats-combine-f32.wgsl.jinja", "inputs": {} }, - "bindings": "gqaMatPastRowStatsCombine", + "shader": "attn-materialized-rowstats-combine-f32.wgsl.jinja", + "bindings": ["scorePartials_2", "rowStats", "params_7"], "dispatch": { - "gridStride": "dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)", - "workgroupSize": "gqaMatRowStatsWg" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)), (gqaMatRowStatsWg)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", + "y": 1, + "z": 1 } }, { "id": "apply", "name": "GroupQueryAttention.PastMaterializedApply", - "source": { - "shader": "attn-materialized-sgmat-f32.wgsl.jinja", - "inputs": { - "phase": "\"apply\"", - "fusedSoftmax": true, - "layout": "\"bsh\"", - "kvLayout": "\"bhsd\"", - "causalRightAlign": true - } + "shader": "attn-materialized-sgmat-f32.wgsl.jinja", + "derive": { + "phase": "\"apply\"", + "fusedSoftmax": true, + "layout": "\"bsh\"", + "kvLayout": "\"bhsd\"", + "causalRightAlign": true }, - "bindings": "gqaMatPastApplyFused", + "bindings": ["scores_2", "value_3", "rowStats_2", "output", "params_8"], "dispatch": { "x": "ceilDiv(gqaHeadDim, gqaMatKeyTile)", "y": "ceilDiv(dim(shapes.queryT, 1), gqaMatQueryTile)", @@ -2560,14 +894,13 @@ }, { "id": "past_kv_bias_materialized_sgmat_f32", + "priority": 36, + "when": ["gqaMatPastBiasOk"], "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": { + "derive": { "qNumHeads": "attrs.num_heads", "kvNumHeads": "attrs.kv_num_heads", "headDim": "gqaHeadDim", @@ -2588,8 +921,7 @@ "packed": "gqaHeadDim", "kvHeads": "attrs.kv_num_heads", "inputScalar": "gqaScalar", - "usesF16": false, - "copyWorkgroupSize": "copyWorkgroupSize" + "usesF16": false }, "intermediates": [ { @@ -2605,58 +937,73 @@ "id": "present", "name": "GroupQueryAttention.Present", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentCopy", + "bindings": ["src_k", "src_v", "present_key_2", "present_value_2", "params_5"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "z": 1 }, "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 - } + "shader": "attn-materialized-sgmat-f32.wgsl.jinja", + "derive": { + "phase": "\"score\"", + "emitRowStats": true, + "layout": "\"bsh\"", + "kvLayout": "\"bhsd\"", + "causalRightAlign": true, + "scoreBias": true }, - "bindings": "gqaMatPastBiasScoreStats", + "bindings": [ + "query", + "key_3", + { "arg": "attentionBiasT", "name": "attention_bias", "elementType": "f32" }, + "scores", + "scorePartials", + { + "name": "params", + "struct": [ + { "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)" } + ] + } + ], "dispatch": { "x": "ceilDiv(dim(shapes.presentKeyT, 2), gqaMatKeyTile)", "y": "ceilDiv(dim(shapes.queryT, 1), gqaMatQueryTile)", "z": "dim(shapes.queryT, 0) * attrs.num_heads" - } + }, + "subgroupCollectivesWidth": 32 }, { "id": "rowstats", "name": "GroupQueryAttention.PastMaterializedRowStats", - "source": { "shader": "attn-materialized-rowstats-combine-f32.wgsl.jinja", "inputs": {} }, - "bindings": "gqaMatPastRowStatsCombine", + "shader": "attn-materialized-rowstats-combine-f32.wgsl.jinja", + "bindings": ["scorePartials_2", "rowStats", "params_7"], "dispatch": { - "gridStride": "dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)", - "workgroupSize": "gqaMatRowStatsWg" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)), (gqaMatRowStatsWg)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", + "y": 1, + "z": 1 } }, { "id": "apply", "name": "GroupQueryAttention.PastMaterializedApply", - "source": { - "shader": "attn-materialized-sgmat-f32.wgsl.jinja", - "inputs": { - "phase": "\"apply\"", - "fusedSoftmax": true, - "layout": "\"bsh\"", - "kvLayout": "\"bhsd\"", - "causalRightAlign": true - } + "shader": "attn-materialized-sgmat-f32.wgsl.jinja", + "derive": { + "phase": "\"apply\"", + "fusedSoftmax": true, + "layout": "\"bsh\"", + "kvLayout": "\"bhsd\"", + "causalRightAlign": true }, - "bindings": "gqaMatPastApplyFused", + "bindings": ["scores_2", "value_3", "rowStats_2", "output", "params_8"], "dispatch": { "x": "ceilDiv(gqaHeadDim, gqaMatKeyTile)", "y": "ceilDiv(dim(shapes.queryT, 1), gqaMatQueryTile)", @@ -2667,14 +1014,13 @@ }, { "id": "past_kv_materialized_sgmat_f16", + "priority": 35, + "when": ["gqaMatPastCoreF16Ok"], "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": { + "derive": { "qNumHeads": "attrs.num_heads", "kvNumHeads": "attrs.kv_num_heads", "headDim": "gqaHeadDim", @@ -2696,7 +1042,6 @@ "kvHeads": "attrs.kv_num_heads", "inputScalar": "gqaScalar", "usesF16": true, - "copyWorkgroupSize": "copyWorkgroupSize", "operandF16": true }, "intermediates": [ @@ -2713,57 +1058,62 @@ "id": "present", "name": "GroupQueryAttention.Present", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentCopy", + "bindings": ["src_k", "src_v", "present_key_2", "present_value_2", "params_5"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "z": 1 }, "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 - } + "shader": "attn-materialized-sgmat-f32.wgsl.jinja", + "derive": { + "phase": "\"score\"", + "emitRowStats": true, + "layout": "\"bsh\"", + "kvLayout": "\"bhsd\"", + "causalRightAlign": true }, - "bindings": "gqaMatPastScoreStatsF16", + "bindings": [ + { "arg": "queryT", "name": "query", "elementType": "f16" }, + "key_4", + "scores", + "scorePartials", + "params_6" + ], "dispatch": { "x": "ceilDiv(dim(shapes.presentKeyT, 2), gqaMatKeyTile)", "y": "ceilDiv(dim(shapes.queryT, 1), gqaMatQueryTile)", "z": "dim(shapes.queryT, 0) * attrs.num_heads" - } + }, + "subgroupCollectivesWidth": 32 }, { "id": "rowstats", "name": "GroupQueryAttention.PastMaterializedRowStats", - "source": { "shader": "attn-materialized-rowstats-combine-f32.wgsl.jinja", "inputs": {} }, - "bindings": "gqaMatPastRowStatsCombine", + "shader": "attn-materialized-rowstats-combine-f32.wgsl.jinja", + "bindings": ["scorePartials_2", "rowStats", "params_7"], "dispatch": { - "gridStride": "dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)", - "workgroupSize": "gqaMatRowStatsWg" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)), (gqaMatRowStatsWg)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", + "y": 1, + "z": 1 } }, { "id": "apply", "name": "GroupQueryAttention.PastMaterializedApply", - "source": { - "shader": "attn-materialized-sgmat-f32.wgsl.jinja", - "inputs": { - "phase": "\"apply\"", - "fusedSoftmax": true, - "layout": "\"bsh\"", - "kvLayout": "\"bhsd\"", - "causalRightAlign": true - } + "shader": "attn-materialized-sgmat-f32.wgsl.jinja", + "derive": { + "phase": "\"apply\"", + "fusedSoftmax": true, + "layout": "\"bsh\"", + "kvLayout": "\"bhsd\"", + "causalRightAlign": true }, - "bindings": "gqaMatPastApplyFusedF16", + "bindings": ["scores_2", "value_4", "rowStats_2", "output_2", "params_8"], "dispatch": { "x": "ceilDiv(gqaHeadDim, gqaMatKeyTile)", "y": "ceilDiv(dim(shapes.queryT, 1), gqaMatQueryTile)", @@ -2774,14 +1124,13 @@ }, { "id": "past_kv_rotary_materialized_sgmat_f32", + "priority": 37, + "when": ["gqaMatPastRotaryOk"], "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": { + "derive": { "qNumHeads": "attrs.num_heads", "kvNumHeads": "attrs.kv_num_heads", "headDim": "gqaHeadDim", @@ -2803,7 +1152,6 @@ "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\"", @@ -2828,10 +1176,11 @@ "id": "present", "name": "GroupQueryAttention.Present", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentCopy", + "bindings": ["src_k", "src_v", "present_key_2", "present_value_2", "params_5"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "z": 1 }, "viewAlias": [{ "input": "src_k", "output": "present_key" }, { "input": "src_v", "output": "present_value" }] }, @@ -2839,56 +1188,61 @@ "id": "qprep", "name": "GroupQueryAttention.RotaryQPrep", "shader": "gqa-qprep.wgsl.jinja", - "bindings": "rotaryQprep", + "bindings": ["query_3", "cos_cache", "sin_cache", "qout", "params_10"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)), (copyWorkgroupSize)), 65535)", + "z": 1 } }, { "id": "scores", "name": "GroupQueryAttention.PastRotaryMaterializedScores", - "source": { - "shader": "attn-materialized-sgmat-f32.wgsl.jinja", - "inputs": { - "phase": "\"score\"", - "emitRowStats": true, - "layout": "\"bsh\"", - "kvLayout": "\"bhsd\"", - "causalRightAlign": true - } + "shader": "attn-materialized-sgmat-f32.wgsl.jinja", + "derive": { + "phase": "\"score\"", + "emitRowStats": true, + "layout": "\"bsh\"", + "kvLayout": "\"bhsd\"", + "causalRightAlign": true }, - "bindings": "gqaMatPastRotaryScoreStats", + "bindings": [ + { "scratch": "qPrep", "name": "query", "buffer": "read-only-storage", "elementType": "f32" }, + "key_3", + "scores", + "scorePartials", + "params_6" + ], "dispatch": { "x": "ceilDiv(dim(shapes.presentKeyT, 2), gqaMatKeyTile)", "y": "ceilDiv(dim(shapes.queryT, 1), gqaMatQueryTile)", "z": "dim(shapes.queryT, 0) * attrs.num_heads" - } + }, + "subgroupCollectivesWidth": 32 }, { "id": "rowstats", "name": "GroupQueryAttention.PastRotaryMaterializedRowStats", - "source": { "shader": "attn-materialized-rowstats-combine-f32.wgsl.jinja", "inputs": {} }, - "bindings": "gqaMatPastRowStatsCombine", + "shader": "attn-materialized-rowstats-combine-f32.wgsl.jinja", + "bindings": ["scorePartials_2", "rowStats", "params_7"], "dispatch": { - "gridStride": "dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)", - "workgroupSize": "gqaMatRowStatsWg" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)), (gqaMatRowStatsWg)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", + "y": 1, + "z": 1 } }, { "id": "apply", "name": "GroupQueryAttention.PastRotaryMaterializedApply", - "source": { - "shader": "attn-materialized-sgmat-f32.wgsl.jinja", - "inputs": { - "phase": "\"apply\"", - "fusedSoftmax": true, - "layout": "\"bsh\"", - "kvLayout": "\"bhsd\"", - "causalRightAlign": true - } + "shader": "attn-materialized-sgmat-f32.wgsl.jinja", + "derive": { + "phase": "\"apply\"", + "fusedSoftmax": true, + "layout": "\"bsh\"", + "kvLayout": "\"bhsd\"", + "causalRightAlign": true }, - "bindings": "gqaMatPastApplyFused", + "bindings": ["scores_2", "value_3", "rowStats_2", "output", "params_8"], "dispatch": { "x": "ceilDiv(gqaHeadDim, gqaMatKeyTile)", "y": "ceilDiv(dim(shapes.queryT, 1), gqaMatQueryTile)", @@ -2899,14 +1253,13 @@ }, { "id": "past_kv_rotary_materialized_sgmat_f16", + "priority": 37, + "when": ["gqaMatPastRotaryF16Ok"], "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": { + "derive": { "qNumHeads": "attrs.num_heads", "kvNumHeads": "attrs.kv_num_heads", "headDim": "gqaHeadDim", @@ -2928,7 +1281,6 @@ "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\"", @@ -2954,10 +1306,11 @@ "id": "present", "name": "GroupQueryAttention.Present", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentCopy", + "bindings": ["src_k", "src_v", "present_key_2", "present_value_2", "params_5"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "z": 1 }, "viewAlias": [{ "input": "src_k", "output": "present_key" }, { "input": "src_v", "output": "present_value" }] }, @@ -2965,57 +1318,68 @@ "id": "qprep", "name": "GroupQueryAttention.RotaryQPrep", "shader": "gqa-qprep.wgsl.jinja", - "bindings": "rotaryQprepF16", + "derive": { "qPrepF16": true }, + "bindings": [ + "query_3", + "cos_cache", + "sin_cache", + { "scratch": "qPrep", "name": "qout", "elementType": "vec4" }, + "params_10" + ], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)", - "workgroupSize": "copyWorkgroupSize" - }, - "constants": { "qPrepF16": true } + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)), (copyWorkgroupSize)), 65535)", + "z": 1 + } }, { "id": "scores", "name": "GroupQueryAttention.PastRotaryMaterializedScores", - "source": { - "shader": "attn-materialized-sgmat-f32.wgsl.jinja", - "inputs": { - "phase": "\"score\"", - "emitRowStats": true, - "layout": "\"bsh\"", - "kvLayout": "\"bhsd\"", - "causalRightAlign": true - } + "shader": "attn-materialized-sgmat-f32.wgsl.jinja", + "derive": { + "phase": "\"score\"", + "emitRowStats": true, + "layout": "\"bsh\"", + "kvLayout": "\"bhsd\"", + "causalRightAlign": true }, - "bindings": "gqaMatPastRotaryScoreStatsF16", + "bindings": [ + { "scratch": "qPrep", "name": "query", "buffer": "read-only-storage", "elementType": "f16" }, + "key_4", + "scores", + "scorePartials", + "params_6" + ], "dispatch": { "x": "ceilDiv(dim(shapes.presentKeyT, 2), gqaMatKeyTile)", "y": "ceilDiv(dim(shapes.queryT, 1), gqaMatQueryTile)", "z": "dim(shapes.queryT, 0) * attrs.num_heads" - } + }, + "subgroupCollectivesWidth": 32 }, { "id": "rowstats", "name": "GroupQueryAttention.PastRotaryMaterializedRowStats", - "source": { "shader": "attn-materialized-rowstats-combine-f32.wgsl.jinja", "inputs": {} }, - "bindings": "gqaMatPastRowStatsCombine", + "shader": "attn-materialized-rowstats-combine-f32.wgsl.jinja", + "bindings": ["scorePartials_2", "rowStats", "params_7"], "dispatch": { - "gridStride": "dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)", - "workgroupSize": "gqaMatRowStatsWg" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)), (gqaMatRowStatsWg)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", + "y": 1, + "z": 1 } }, { "id": "apply", "name": "GroupQueryAttention.PastRotaryMaterializedApply", - "source": { - "shader": "attn-materialized-sgmat-f32.wgsl.jinja", - "inputs": { - "phase": "\"apply\"", - "fusedSoftmax": true, - "layout": "\"bsh\"", - "kvLayout": "\"bhsd\"", - "causalRightAlign": true - } + "shader": "attn-materialized-sgmat-f32.wgsl.jinja", + "derive": { + "phase": "\"apply\"", + "fusedSoftmax": true, + "layout": "\"bsh\"", + "kvLayout": "\"bhsd\"", + "causalRightAlign": true }, - "bindings": "gqaMatPastApplyFusedF16", + "bindings": ["scores_2", "value_4", "rowStats_2", "output_2", "params_8"], "dispatch": { "x": "ceilDiv(gqaHeadDim, gqaMatKeyTile)", "y": "ceilDiv(dim(shapes.queryT, 1), gqaMatQueryTile)", @@ -3026,14 +1390,13 @@ }, { "id": "new_kv_past_materialized_sgmat_f32", + "priority": 35, + "when": ["gqaMatNewPastOk"], "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": { + "derive": { "qNumHeads": "attrs.num_heads", "kvNumHeads": "attrs.kv_num_heads", "headDim": "gqaHeadDim", @@ -3054,8 +1417,7 @@ "packed": "gqaHeadDim", "kvHeads": "attrs.kv_num_heads", "inputScalar": "gqaScalar", - "usesF16": false, - "copyWorkgroupSize": "copyWorkgroupSize" + "usesF16": false }, "intermediates": [ { @@ -3071,56 +1433,55 @@ "id": "present", "name": "GroupQueryAttention.Merge", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentMerge", + "bindings": ["past_k", "past_v", "new_k", "new_v", "present_key_2", "present_value_2", "params_11"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)), (copyWorkgroupSize)), 65535)", + "z": 1 } }, { "id": "scores", "name": "GroupQueryAttention.NewPastMaterializedScores", - "source": { - "shader": "attn-materialized-sgmat-f32.wgsl.jinja", - "inputs": { - "phase": "\"score\"", - "emitRowStats": true, - "layout": "\"bsh\"", - "kvLayout": "\"bhsd\"", - "causalRightAlign": true - } + "shader": "attn-materialized-sgmat-f32.wgsl.jinja", + "derive": { + "phase": "\"score\"", + "emitRowStats": true, + "layout": "\"bsh\"", + "kvLayout": "\"bhsd\"", + "causalRightAlign": true }, - "bindings": "gqaMatPastScoreStats", + "bindings": ["query", "key_3", "scores", "scorePartials", "params_6"], "dispatch": { "x": "ceilDiv(dim(shapes.presentKeyT, 2), gqaMatKeyTile)", "y": "ceilDiv(dim(shapes.queryT, 1), gqaMatQueryTile)", "z": "dim(shapes.queryT, 0) * attrs.num_heads" - } + }, + "subgroupCollectivesWidth": 32 }, { "id": "rowstats", "name": "GroupQueryAttention.NewPastMaterializedRowStats", - "source": { "shader": "attn-materialized-rowstats-combine-f32.wgsl.jinja", "inputs": {} }, - "bindings": "gqaMatPastRowStatsCombine", + "shader": "attn-materialized-rowstats-combine-f32.wgsl.jinja", + "bindings": ["scorePartials_2", "rowStats", "params_7"], "dispatch": { - "gridStride": "dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)", - "workgroupSize": "gqaMatRowStatsWg" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)), (gqaMatRowStatsWg)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", + "y": 1, + "z": 1 } }, { "id": "apply", "name": "GroupQueryAttention.NewPastMaterializedApply", - "source": { - "shader": "attn-materialized-sgmat-f32.wgsl.jinja", - "inputs": { - "phase": "\"apply\"", - "fusedSoftmax": true, - "layout": "\"bsh\"", - "kvLayout": "\"bhsd\"", - "causalRightAlign": true - } + "shader": "attn-materialized-sgmat-f32.wgsl.jinja", + "derive": { + "phase": "\"apply\"", + "fusedSoftmax": true, + "layout": "\"bsh\"", + "kvLayout": "\"bhsd\"", + "causalRightAlign": true }, - "bindings": "gqaMatPastApplyFused", + "bindings": ["scores_2", "value_3", "rowStats_2", "output", "params_8"], "dispatch": { "x": "ceilDiv(gqaHeadDim, gqaMatKeyTile)", "y": "ceilDiv(dim(shapes.queryT, 1), gqaMatQueryTile)", @@ -3131,14 +1492,13 @@ }, { "id": "window_shift_materialized_sgmat_f32", + "priority": 36, + "when": ["gqaMatWindowOk"], "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": { + "derive": { "qNumHeads": "attrs.num_heads", "kvNumHeads": "attrs.kv_num_heads", "headDim": "gqaHeadDim", @@ -3160,8 +1520,7 @@ "kvHeads": "attrs.kv_num_heads", "inputScalar": "gqaScalar", "zeroScalar": "gqaScalar", - "usesF16": false, - "copyWorkgroupSize": "copyWorkgroupSize" + "usesF16": false }, "intermediates": [ { @@ -3177,60 +1536,210 @@ "id": "present", "name": "GroupQueryAttention.WindowShift", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentMergeShare", + "bindings": ["past_k", "past_v", "new_k", "new_v", "present_key_2", "present_value_2", "seqlens_k", "params_12"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)), (copyWorkgroupSize)), 65535)", + "z": 1 } }, { "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 - } + "shader": "attn-materialized-sgmat-f32.wgsl.jinja", + "derive": { + "phase": "\"score\"", + "emitRowStats": true, + "layout": "\"bsh\"", + "kvLayout": "\"bhsd\"", + "causalRightAlign": true, + "scoreWindow": true, + "useSeqlens": true }, - "bindings": "gqaMatWindowScoreStats", + "bindings": [ + "query", + "key_3", + "scores", + "scorePartials", + "seqlens_k", + { + "name": "params", + "struct": [ + { "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" } + ] + } + ], "dispatch": { "x": "ceilDiv(dim(shapes.presentKeyT, 2), gqaMatKeyTile)", "y": "ceilDiv(dim(shapes.queryT, 1), gqaMatQueryTile)", "z": "dim(shapes.queryT, 0) * attrs.num_heads" - } + }, + "subgroupCollectivesWidth": 32 }, { "id": "rowstats", "name": "GroupQueryAttention.WindowMaterializedRowStats", - "source": { "shader": "attn-materialized-rowstats-combine-f32.wgsl.jinja", "inputs": {} }, - "bindings": "gqaMatPastRowStatsCombine", + "shader": "attn-materialized-rowstats-combine-f32.wgsl.jinja", + "bindings": ["scorePartials_2", "rowStats", "params_7"], "dispatch": { - "gridStride": "dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)", - "workgroupSize": "gqaMatRowStatsWg" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)), (gqaMatRowStatsWg)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", + "y": 1, + "z": 1 } }, { "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 + "shader": "attn-materialized-sgmat-f32.wgsl.jinja", + "derive": { + "phase": "\"apply\"", + "fusedSoftmax": true, + "layout": "\"bsh\"", + "kvLayout": "\"bhsd\"", + "causalRightAlign": true, + "scoreWindow": true, + "useSeqlens": true + }, + "bindings": [ + "scores_2", + "value_3", + "rowStats_2", + "output", + "seqlens_k", + { + "name": "params", + "struct": [ + { "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" } + ] } + ], + "dispatch": { + "x": "ceilDiv(gqaHeadDim, gqaMatKeyTile)", + "y": "ceilDiv(dim(shapes.queryT, 1), gqaMatQueryTile)", + "z": "dim(shapes.queryT, 0) * attrs.num_heads" + } + } + ] + }, + { + "id": "share_append_materialized_sgmat_f32", + "priority": 36, + "when": ["gqaMatShareAppendOk"], + "requires": { + "features": ["subgroups", "chromium-experimental-subgroup-matrix"], + "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] + }, + "derive": { + "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_share\"", + "useSeqlens": true, + "packed": "gqaHeadDim", + "kvHeads": "attrs.kv_num_heads", + "inputScalar": "gqaScalar", + "usesF16": false + }, + "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_retain", + "name": "GroupQueryAttention.MergeShareRetain", + "shader": "gqa-present.wgsl.jinja", + "derive": { "shareRegion": "\"retain\"" }, + "bindings": ["past_k", "past_v", "present_key_2", "present_value_2", "seqlens_k", "params_12"], + "dispatch": { + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)), (copyWorkgroupSize)), 65535)", + "z": 1 + }, + "viewAlias": [ + { "input": "past_k", "output": "present_key" }, + { "input": "past_v", "output": "present_value" } + ] + }, + { + "id": "present_append", + "name": "GroupQueryAttention.MergeShareAppend", + "shader": "gqa-present.wgsl.jinja", + "derive": { "shareRegion": "\"append\"" }, + "bindings": ["new_k", "new_v", "present_key_2", "present_value_2", "seqlens_k", "params_15"], + "dispatch": { + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)), (copyWorkgroupSize)), 65535)", + "z": 1 + } + }, + { + "id": "scores", + "name": "GroupQueryAttention.ShareAppendMaterializedScores", + "shader": "attn-materialized-sgmat-f32.wgsl.jinja", + "derive": { + "phase": "\"score\"", + "emitRowStats": true, + "layout": "\"bsh\"", + "kvLayout": "\"bhsd\"", + "causalRightAlign": true, + "useSeqlens": true }, - "bindings": "gqaMatWindowApplyFused", + "bindings": ["query", "key_3", "scores", "scorePartials", "seqlens_k", "params_6"], + "dispatch": { + "x": "ceilDiv(dim(shapes.presentKeyT, 2), gqaMatKeyTile)", + "y": "ceilDiv(dim(shapes.queryT, 1), gqaMatQueryTile)", + "z": "dim(shapes.queryT, 0) * attrs.num_heads" + }, + "subgroupCollectivesWidth": 32 + }, + { + "id": "rowstats", + "name": "GroupQueryAttention.ShareAppendMaterializedRowStats", + "shader": "attn-materialized-rowstats-combine-f32.wgsl.jinja", + "bindings": ["scorePartials_2", "rowStats", "params_7"], + "dispatch": { + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)), (gqaMatRowStatsWg)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", + "y": 1, + "z": 1 + } + }, + { + "id": "apply", + "name": "GroupQueryAttention.ShareAppendMaterializedApply", + "shader": "attn-materialized-sgmat-f32.wgsl.jinja", + "derive": { + "phase": "\"apply\"", + "fusedSoftmax": true, + "layout": "\"bsh\"", + "kvLayout": "\"bhsd\"", + "causalRightAlign": true, + "useSeqlens": true + }, + "bindings": ["scores_2", "value_3", "rowStats_2", "output", "seqlens_k", "params_8"], "dispatch": { "x": "ceilDiv(gqaHeadDim, gqaMatKeyTile)", "y": "ceilDiv(dim(shapes.queryT, 1), gqaMatQueryTile)", @@ -3241,10 +1750,9 @@ }, { "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": { + "derive": { "useSeqlens": true, "headDim": "gqaHeadDim", "cooperative": "gqaCoop", @@ -3262,11 +1770,12 @@ "id": "present_retain", "name": "GroupQueryAttention.MergeShareRetain", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentMergeShareRetain", - "constants": { "shareRegion": "\"retain\"" }, + "derive": { "shareRegion": "\"retain\"" }, + "bindings": ["past_k", "past_v", "present_key_2", "present_value_2", "seqlens_k", "params_12"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)), (copyWorkgroupSize)), 65535)", + "z": 1 }, "viewAlias": [ { "input": "past_k", "output": "present_key" }, @@ -3277,19 +1786,20 @@ "id": "present_append", "name": "GroupQueryAttention.MergeShareAppend", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentMergeShareAppend", - "constants": { "shareRegion": "\"append\"" }, + "derive": { "shareRegion": "\"append\"" }, + "bindings": ["new_k", "new_v", "present_key_2", "present_value_2", "seqlens_k", "params_15"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)), (copyWorkgroupSize)), 65535)", + "z": 1 } }, { "id": "attention", "name": "GroupQueryAttention.Attn", "shader": "gqa-attention.wgsl.jinja", - "bindings": "mergedSeqlensAttention", - "dispatch": { "workgroups": "constants.dispatchUnits" } + "bindings": ["query_3", "kcache", "vcache", "output_3", "seqlens_k", "params_16"], + "dispatch": { "x": "min(dispatchUnits, 65535)", "y": "ceilDiv(dispatchUnits, 65535)", "z": 1 } } ] }, @@ -3297,7 +1807,7 @@ "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": { + "derive": { "useSeqlens": true, "headDim": "gqaHeadDim", "cooperative": "gqaCoop", @@ -3316,11 +1826,12 @@ "id": "present_retain", "name": "GroupQueryAttention.MergeShareRetain", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentMergeShareRetain", - "constants": { "shareRegion": "\"retain\"" }, + "derive": { "shareRegion": "\"retain\"" }, + "bindings": ["past_k", "past_v", "present_key_2", "present_value_2", "seqlens_k", "params_12"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)), (copyWorkgroupSize)), 65535)", + "z": 1 }, "viewAlias": [ { "input": "past_k", "output": "present_key" }, @@ -3331,19 +1842,42 @@ "id": "present_append", "name": "GroupQueryAttention.MergeShareAppend", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentMergeShareAppend", - "constants": { "shareRegion": "\"append\"" }, + "derive": { "shareRegion": "\"append\"" }, + "bindings": ["new_k", "new_v", "present_key_2", "present_value_2", "seqlens_k", "params_15"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)), (copyWorkgroupSize)), 65535)", + "z": 1 } }, { "id": "attention", "name": "GroupQueryAttention.Attn", "shader": "gqa-attention.wgsl.jinja", - "bindings": "mergedSeqlensHeadSinkAttention", - "dispatch": { "workgroups": "constants.dispatchUnits" } + "bindings": [ + "query_3", + "kcache", + "vcache", + "head_sink", + "output_3", + "seqlens_k", + { + "name": "params", + "struct": [ + { "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" } + ] + } + ], + "dispatch": { "x": "min(dispatchUnits, 65535)", "y": "ceilDiv(dispatchUnits, 65535)", "z": 1 } } ] }, @@ -3351,7 +1885,7 @@ "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": { + "derive": { "useSeqlens": true, "headDim": "gqaHeadDim", "cooperative": "gqaCoop", @@ -3371,11 +1905,12 @@ "id": "present_retain", "name": "GroupQueryAttention.MergeShareRetain", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentMergeShareRetain", - "constants": { "shareRegion": "\"retain\"" }, + "derive": { "shareRegion": "\"retain\"" }, + "bindings": ["past_k", "past_v", "present_key_2", "present_value_2", "seqlens_k", "params_12"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)), (copyWorkgroupSize)), 65535)", + "z": 1 }, "viewAlias": [ { "input": "past_k", "output": "present_key" }, @@ -3386,28 +1921,28 @@ "id": "present_append", "name": "GroupQueryAttention.MergeShareAppend", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentMergeShareRotaryAppend", - "constants": { "shareRegion": "\"append\"" }, + "derive": { "shareRegion": "\"append\"" }, + "bindings": ["new_k", "new_v", "cos_cache_2", "sin_cache_2", "present_key_2", "present_value_2", "seqlens_k", "params_15"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)), (copyWorkgroupSize)), 65535)", + "z": 1 } }, { "id": "attention", "name": "GroupQueryAttention.Attn", "shader": "gqa-attention.wgsl.jinja", - "bindings": "mergedSeqlensRotaryAttention", - "dispatch": { "workgroups": "constants.dispatchUnits" } + "bindings": ["query_3", "kcache", "vcache", "cos_cache_2", "sin_cache_2", "output_3", "seqlens_k", "params_16"], + "dispatch": { "x": "min(dispatchUnits, 65535)", "y": "ceilDiv(dispatchUnits, 65535)", "z": 1 } } ] }, { "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": { + "when": ["plainAttentionOptions", "flashShapeOk", "qkvTiledStorageOk", "headDimPlan % 4 == 0", "headDimPlan <= 128", "dim(shapes.queryT, 1) >= tunables.QKV_PREFILL_MIN_QUERY_TOKENS", "attrs.local_window_size <= 0"], + "derive": { "scalar": "qkvScalar", "inputScalar": "qkvScalar", "outputScalar": "qkvScalar", @@ -3417,40 +1952,67 @@ "presentElement": "qkvCacheScalar", "usesF16": "qkvUsesF16", "blockM": 32, - "vHeadCap": "dim(shapes.value, 2) / attrs.kv_num_heads", + "vHeadCap": "dim(shapes.valueT, 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" + "headDim": "dim(shapes.queryT, 2) / attrs.num_heads", + "qHidden": "dim(shapes.queryT, 2)", + "kvHidden": "dim(shapes.keyT, 2)" }, "passes": [ { "id": "attention", "name": "GroupQueryAttention.PrefillTiledNoSg", - "source": { "shader": "attention-rank4-tiled.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, - "constants": { "usesF16": "tensorDtypes.query == \"float16\"" }, - "bindings": "qkvTiledAttention", + "shader": "attention-rank4-tiled.wgsl.jinja", + "derive": { "usesF16": "tensorDtypes.queryT == \"float16\"", "layout": "\"bsh\"" }, + "bindings": [ + { "arg": "queryT", "name": "q", "elementType": "$scalar" }, + { "arg": "keyT", "name": "k", "elementType": "$scalar" }, + { "arg": "valueT", "name": "v", "elementType": "$scalar" }, + { "arg": "outputT", "name": "y", "elementType": "$scalar" }, + { + "name": "params", + "struct": [ + { "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.queryT, 1)" }, + { "name": "kvSeq", "type": "u32", "value": "dim(shapes.keyT, 1)" }, + { "name": "headSize", "type": "u32", "value": "dim(shapes.queryT, 2) / attrs.num_heads" }, + { "name": "vHeadSize", "type": "u32", "value": "dim(shapes.valueT, 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.queryT, 2)" }, + { "name": "kvHidden", "type": "u32", "value": "dim(shapes.keyT, 2)" }, + { "name": "vHidden", "type": "u32", "value": "dim(shapes.valueT, 2)" } + ] + } + ], "dispatch": { - "threads": "dim(shapes.outputT, 0) * attrs.num_heads * ceilDiv(dim(shapes.outputT, 1), 32) * 32", - "workgroupSize": "32" + "x": "min(ceilDiv((dim(shapes.outputT, 0) * attrs.num_heads * ceilDiv(dim(shapes.outputT, 1), 32) * 32), (32)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.outputT, 0) * attrs.num_heads * ceilDiv(dim(shapes.outputT, 1), 32) * 32), (32)), 65535)", + "z": 1 } }, { "id": "present", "name": "GroupQueryAttention.PresentTiledNoSg", "shader": "gqa-present.wgsl.jinja", - "bindings": "qkvPresent", - "dispatch": { "threads": "numel(shapes.presentKeyT)", "workgroupSize": "constants.copyWorkgroupSize" } + "bindings": ["key_2", "value_2", "present_key", "present_value", "params"], + "dispatch": { + "x": "min(ceilDiv((numel(shapes.presentKeyT)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((numel(shapes.presentKeyT)), (copyWorkgroupSize)), 65535)", + "z": 1 + } } ] }, { "id": "qkv_present_flash", "priority": 20, + "when": ["flashShapeOk"], "requires": { "features": ["subgroups"] }, - "when": "flashShapeOk", - "constants": { + "derive": { "headsFromParams": false, "hasBias": false, "hasCausal": false, @@ -3465,12 +2027,11 @@ "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", + "headDim": "dim(shapes.queryT, 2) / attrs.num_heads", + "headDimV4": "(dim(shapes.queryT, 2) / attrs.num_heads) / 4", + "qHiddenV4": "dim(shapes.queryT, 2) / 4", + "kvHidden": "dim(shapes.keyT, 2)", + "kvHiddenV4": "dim(shapes.keyT, 2) / 4", "presentVec4": true, "presentElem": "qkvCacheVec4", "hasWindow": "attrs.local_window_size > 0" @@ -3479,25 +2040,30 @@ { "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)" } + "shader": "attn-flash-online.wgsl.jinja", + "derive": { "usesF16": "tensorDtypes.queryT == \"float16\"", "layout": "\"bsh\"" }, + "bindings": ["query_6", "key_2", "value_2", "output_4", "params_19"], + "dispatch": { "x": "dim(shapes.queryT, 1)", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }, + "subgroupCollectivesWidth": "portable" }, { "id": "present", "name": "GroupQueryAttention.Present", "shader": "gqa-present.wgsl.jinja", - "bindings": "qkvPresent", - "dispatch": { "threads": "numel(shapes.presentKeyT) / 4", "workgroupSize": "constants.copyWorkgroupSize" } + "bindings": ["key_2", "value_2", "present_key", "present_value", "params"], + "dispatch": { + "x": "min(ceilDiv((numel(shapes.presentKeyT) / 4), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((numel(shapes.presentKeyT) / 4), (copyWorkgroupSize)), 65535)", + "z": 1 + } } ] }, { "id": "qkv_present", "priority": 0, - "when": "fallbackShapeOk", - "constants": { + "when": ["fallbackShapeOk"], + "derive": { "headsFromParams": false, "hasBias": false, "hasCausal": false, @@ -3512,32 +2078,35 @@ "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", + "headDim": "dim(shapes.queryT, 2) / attrs.num_heads", + "qHidden": "dim(shapes.queryT, 2)", + "kvHidden": "dim(shapes.keyT, 2)", + "workgroupSize": "min(256, pow2ceil(dim(shapes.queryT, 2) / attrs.num_heads))", "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", + "shader": "attn-online-scalar.wgsl.jinja", + "derive": { "usesF16": "tensorDtypes.queryT == \"float16\"", "layout": "\"bsh\"" }, + "bindings": ["query_6", "key_2", "value_2", "output_4", "params_19"], "dispatch": { - "x": "min(dim(shapes.query, 1), device.limits.maxComputeWorkgroupsPerDimension)", + "x": "min(dim(shapes.queryT, 1), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))", "y": "attrs.num_heads", - "z": "dim(shapes.query, 0)" + "z": "dim(shapes.queryT, 0)" } }, { "id": "present", "name": "GroupQueryAttention.Present", "shader": "gqa-present.wgsl.jinja", - "bindings": "qkvPresent", - "dispatch": { "threads": "numel(shapes.presentKeyT)", "workgroupSize": "constants.copyWorkgroupSize" } + "bindings": ["key_2", "value_2", "present_key", "present_value", "params"], + "dispatch": { + "x": "min(ceilDiv((numel(shapes.presentKeyT)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((numel(shapes.presentKeyT)), (copyWorkgroupSize)), 65535)", + "z": 1 + } } ] }, @@ -3545,7 +2114,7 @@ "id": "quant_int8", "priority": 35, "when": ["quantizedPromptOk", "gqaScalarDispatchOk", "attrs.kv_cache_bit_width == 8"], - "constants": { + "derive": { "headDim": "gqaHeadDim", "cooperative": "gqaCoop", "dispatchUnits": "gqaDispatchUnits", @@ -3565,18 +2134,53 @@ "id": "present", "name": "GroupQueryAttention.Present", "shader": "gqa-present.wgsl.jinja", - "bindings": "quantPresent", + "bindings": ["src_k_2", "src_v_2", "k_scale", "v_scale", "present_key_3", "present_value_3", "params_20"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)), (copyWorkgroupSize)), 65535)", + "z": 1 } }, { "id": "attention", "name": "GroupQueryAttention.Attn", "shader": "gqa-attention.wgsl.jinja", - "bindings": "quantAttention", - "dispatch": { "workgroups": "constants.dispatchUnits" } + "bindings": [ + "query", + { + "arg": "presentKeyT", + "name": "kcache", + "buffer": "read-only-storage", + "elementType": "$quantCacheElement" + }, + { + "arg": "presentValueT", + "name": "vcache", + "buffer": "read-only-storage", + "elementType": "$quantCacheElement" + }, + "k_scale", + "v_scale", + "output", + { + "name": "params", + "struct": [ + { "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" } + ] + } + ], + "dispatch": { "x": "min(dispatchUnits, 65535)", "y": "ceilDiv(dispatchUnits, 65535)", "z": 1 } } ] }, @@ -3584,7 +2188,7 @@ "id": "quant_int4", "priority": 35, "when": ["quantizedPromptOk", "gqaScalarDispatchOk", "attrs.kv_cache_bit_width == 4"], - "constants": { + "derive": { "headDim": "gqaHeadDim", "cooperative": "gqaCoop", "dispatchUnits": "gqaDispatchUnits", @@ -3604,27 +2208,61 @@ "id": "present", "name": "GroupQueryAttention.Present", "shader": "gqa-present.wgsl.jinja", - "bindings": "quantPresent", + "bindings": ["src_k_2", "src_v_2", "k_scale", "v_scale", "present_key_3", "present_value_3", "params_20"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)), (copyWorkgroupSize)), 65535)", + "z": 1 } }, { "id": "attention", "name": "GroupQueryAttention.Attn", "shader": "gqa-attention.wgsl.jinja", - "bindings": "quantAttention", - "dispatch": { "workgroups": "constants.dispatchUnits" } + "bindings": [ + "query", + { + "arg": "presentKeyT", + "name": "kcache", + "buffer": "read-only-storage", + "elementType": "$quantCacheElement" + }, + { + "arg": "presentValueT", + "name": "vcache", + "buffer": "read-only-storage", + "elementType": "$quantCacheElement" + }, + "k_scale", + "v_scale", + "output", + { + "name": "params", + "struct": [ + { "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" } + ] + } + ], + "dispatch": { "x": "min(dispatchUnits, 65535)", "y": "ceilDiv(dispatchUnits, 65535)", "z": 1 } } ] }, { "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": { + "derive": { "headDim": "gqaHeadDim", "headDimV4": "gqaHeadDim / 4", "qNumHeads": "attrs.num_heads", @@ -3661,46 +2299,95 @@ { "id": "copy_past", "name": "GroupQueryAttention.QuantInt8CopyPast", - "source": { "shader": "gqa-present.wgsl.jinja", "inputs": { "mode": "\"copy\"" } }, - "bindings": "presentCopy", + "shader": "gqa-present.wgsl.jinja", + "derive": { "modeSpec": "\"copy\"" }, + "bindings": ["src_k", "src_v", "present_key_2", "present_value_2", "params_5"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "z": 1 }, "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", + "shader": "gqa-present.wgsl.jinja", + "derive": { "modeSpec": "\"append_quant\"" }, + "bindings": [ + "src_k_2", + "src_v_2", + "k_scale", + "v_scale", + "seqlens_k", + { "arg": "presentKeyT", "name": "present_key", "elementType": "i32" }, + { "arg": "presentValueT", "name": "present_value", "elementType": "i32" }, + { + "name": "params", + "struct": [ + { + "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" } + ] + } + ], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)), (copyWorkgroupSize)), 65535)", + "z": 1 } }, { "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)" } + "shader": "attn-flash-decode-splitk.wgsl.jinja", + "derive": { "layout": "\"bhsd\"" }, + "bindings": [ + { "arg": "queryT", "name": "query", "elementType": "vec4" }, + { "arg": "presentKeyT", "name": "key", "buffer": "read-only-storage", "elementType": "vec4" }, + { "arg": "presentValueT", "name": "value", "buffer": "read-only-storage", "elementType": "vec4" }, + "k_scale", + "v_scale", + "seqlens_k", + "partial_out", + "partial_stats", + { + "name": "params", + "struct": [ + { "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" } + ] + } + ], + "dispatch": { "x": "cachedNumSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }, + "subgroupCollectivesWidth": "portable" }, { "id": "merge", "name": "GroupQueryAttention.QuantInt8DecodeSplitKMerge", - "source": { "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", "inputs": { "layout": "\"bhsd\"" } }, - "bindings": "splitMergeNoParams", + "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", + "derive": { "layout": "\"bhsd\"" }, + "bindings": ["partial_out_2", "partial_stats_2", "output_5"], "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": { + "derive": { "combineSubgroups": "\"\" == \"\"", "scalar": "qkvScalar", "inputVec4": "qkvInputVec4", @@ -3711,13 +2398,12 @@ "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", + "headDim": "dim(shapes.queryT, 2) / attrs.num_heads", + "headDimV4": "(dim(shapes.queryT, 2) / attrs.num_heads) / 4", + "qHiddenV4": "dim(shapes.queryT, 2) / 4", + "kvHidden": "dim(shapes.keyT, 2)", + "kvHiddenV4": "dim(shapes.keyT, 2) / 4", "numSplits": "qkvNumSplits", - "copyWorkgroupSize": "copyWorkgroupSize", "presentVec4": true, "usesF16": "qkvUsesF16", "presentElem": "qkvCacheVec4", @@ -3728,47 +2414,69 @@ { "id": "partialOut", "dtype": "float32", - "shape": "[dim(shapes.query, 0) * attrs.num_heads * qkvNumSplits * (dim(shapes.query, 2) / attrs.num_heads)]" + "shape": "[dim(shapes.queryT, 0) * attrs.num_heads * qkvNumSplits * (dim(shapes.queryT, 2) / attrs.num_heads)]" }, { "id": "partialStats", "dtype": "float32", - "shape": "[2 * dim(shapes.query, 0) * attrs.num_heads * qkvNumSplits]" + "shape": "[2 * dim(shapes.queryT, 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)" } + "shader": "attn-flash-decode-splitk.wgsl.jinja", + "derive": { "usesF16": "tensorDtypes.queryT == \"float16\"", "layout": "\"bsh\"" }, + "bindings": [ + "query_8", + { "arg": "keyT", "name": "key", "elementType": "$inputVec4" }, + { "arg": "valueT", "name": "value", "elementType": "$inputVec4" }, + "partial_out", + "partial_stats", + { + "name": "params", + "struct": [ + { "name": "kvSeq", "type": "u32", "value": "dim(shapes.keyT, 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" + } + ] + } + ], + "dispatch": { "x": "qkvNumSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }, + "subgroupCollectivesWidth": "portable" }, { "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)" } + "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", + "derive": { "usesF16": "tensorDtypes.queryT == \"float16\"", "layout": "\"bsh\"" }, + "bindings": ["partial_out_2", "partial_stats_2", "output_5"], + "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } }, { "id": "present", "name": "GroupQueryAttention.Present", "shader": "gqa-present.wgsl.jinja", - "bindings": "qkvPresent", - "dispatch": { "threads": "numel(shapes.presentKeyT) / 4", "workgroupSize": "constants.copyWorkgroupSize" } + "bindings": ["key_2", "value_2", "present_key", "present_value", "params"], + "dispatch": { + "x": "min(ceilDiv((numel(shapes.presentKeyT) / 4), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((numel(shapes.presentKeyT) / 4), (copyWorkgroupSize)), 65535)", + "z": 1 + } } ] }, { "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, + "when": ["flashShapeOk", "qkvClusterStorageOk if \"\" == \"\" else qkvNoSgClusterStorageOk", "qkvPrefillClusterWorkgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup", "qkvPrefillClusterWorkgroupSize <= device.limits.maxComputeWorkgroupSizeX", "headDimPlan % (4 * qkvPrefillLanesPerQuery) == 0", "qkvF32ClusterRegisterGeometry if \"\" == \"\" else true", "dim(shapes.queryT, 1) >= tunables.QKV_PREFILL_MIN_QUERY_TOKENS", "subgroupQkvPrefillCluster if \"\" == \"\" else true"], "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": { + "derive": { "hasBias": false, "hasCausal": false, "headsFromParams": false, @@ -3783,12 +2491,11 @@ "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", + "headDim": "dim(shapes.queryT, 2) / attrs.num_heads", + "headDimV4": "(dim(shapes.queryT, 2) / attrs.num_heads) / 4", + "qHiddenV4": "dim(shapes.queryT, 2) / 4", + "kvHidden": "dim(shapes.keyT, 2)", + "kvHiddenV4": "dim(shapes.keyT, 2) / 4", "presentVec4": true, "TILE_Q": "qkvPrefillQueryTile", "TILE_K": "qkvClusterTileK", @@ -3802,21 +2509,26 @@ { "id": "attention", "name": "GroupQueryAttention.Flash", - "source": { "shader": "attn-flash-prefill-cluster.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, - "constants": { "usesF16": "tensorDtypes.query == \"float16\"" }, - "bindings": "qkvAttention", + "shader": "attn-flash-prefill-cluster.wgsl.jinja", + "derive": { "usesF16": "tensorDtypes.queryT == \"float16\"", "layout": "\"bsh\"" }, + "bindings": ["query_6", "key_2", "value_2", "output_4", "params_19"], "dispatch": { - "x": "ceilDiv(dim(shapes.query, 1), constants.TILE_Q)", + "x": "ceilDiv(dim(shapes.queryT, 1), TILE_Q)", "y": "attrs.num_heads", - "z": "dim(shapes.query, 0)" - } + "z": "dim(shapes.queryT, 0)" + }, + "subgroupCollectivesWidth": "portable" }, { "id": "present", "name": "GroupQueryAttention.Present", "shader": "gqa-present.wgsl.jinja", - "bindings": "qkvPresent", - "dispatch": { "threads": "numel(shapes.presentKeyT) / 4", "workgroupSize": "constants.copyWorkgroupSize" } + "bindings": ["key_2", "value_2", "present_key", "present_value", "params"], + "dispatch": { + "x": "min(ceilDiv((numel(shapes.presentKeyT) / 4), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((numel(shapes.presentKeyT) / 4), (copyWorkgroupSize)), 65535)", + "z": 1 + } } ] }, @@ -3824,7 +2536,7 @@ "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": { + "derive": { "scalar": "gqaScalar", "usesF16": "tensorDtypes.queryT == \"float16\"", "hasCausal": true, @@ -3855,26 +2567,26 @@ "id": "present", "name": "GroupQueryAttention.Present", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentCopy", + "bindings": ["src_k", "src_v", "present_key_2", "present_value_2", "params_5"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "z": 1 }, "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", + "shader": "attn-flash-prefill-cluster.wgsl.jinja", + "derive": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" }, + "bindings": ["query_8", "key_7", "value_7", "attn_mask", "output_5", "params_25"], "dispatch": { - "x": "ceilDiv(dim(shapes.queryT, 1), constants.TILE_Q)", + "x": "ceilDiv(dim(shapes.queryT, 1), TILE_Q)", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" - } + }, + "subgroupCollectivesWidth": "portable" } ] }, @@ -3882,7 +2594,7 @@ "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": { + "derive": { "scalar": "gqaScalar", "usesF16": "tensorDtypes.queryT == \"float16\"", "hasCausal": true, @@ -3927,10 +2639,11 @@ "id": "present", "name": "GroupQueryAttention.Present", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentCopy", + "bindings": ["src_k", "src_v", "present_key_2", "present_value_2", "params_5"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "z": 1 }, "viewAlias": [{ "input": "src_k", "output": "present_key" }, { "input": "src_v", "output": "present_value" }] }, @@ -3938,34 +2651,39 @@ "id": "qprep", "name": "GroupQueryAttention.QPrep", "shader": "gqa-qprep.wgsl.jinja", - "bindings": "qnormRotaryQprep", + "bindings": ["query_3", "cos_cache", "sin_cache", "q_norm_weight", "qout", "params_10"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)), (copyWorkgroupSize)), 65535)", + "z": 1 } }, { "id": "attention", "name": "GroupQueryAttention.FlashPrefill", - "source": { - "shader": "attn-flash-prefill-cluster.wgsl.jinja", - "inputs": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" } - }, - "bindings": "rotaryFlashAttention", + "shader": "attn-flash-prefill-cluster.wgsl.jinja", + "derive": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" }, + "bindings": [ + { "scratch": "qPrep", "name": "query", "buffer": "read-only-storage", "elementType": "vec4" }, + "key_7", + "value_7", + "output_5", + "params_26" + ], "dispatch": { - "x": "ceilDiv(dim(shapes.queryT, 1), constants.TILE_Q)", + "x": "ceilDiv(dim(shapes.queryT, 1), TILE_Q)", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" - } + }, + "subgroupCollectivesWidth": "portable" } ] }, { "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": { + "derive": { "headDim": "gqaHeadDim", "headDimV4": "gqaHeadDim / 4", "qNumHeads": "attrs.num_heads", @@ -4002,46 +2720,94 @@ { "id": "copy_past", "name": "GroupQueryAttention.QuantInt8CopyPast", - "source": { "shader": "gqa-present.wgsl.jinja", "inputs": { "mode": "\"copy\"" } }, - "bindings": "presentCopy", + "shader": "gqa-present.wgsl.jinja", + "derive": { "modeSpec": "\"copy\"" }, + "bindings": ["src_k", "src_v", "present_key_2", "present_value_2", "params_5"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "z": 1 }, "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", + "shader": "gqa-present.wgsl.jinja", + "derive": { "modeSpec": "\"append_quant\"" }, + "bindings": [ + "src_k_2", + "src_v_2", + "k_scale", + "v_scale", + "seqlens_k", + { "arg": "presentKeyT", "name": "present_key", "elementType": "i32" }, + { "arg": "presentValueT", "name": "present_value", "elementType": "i32" }, + { + "name": "params", + "struct": [ + { + "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" } + ] + } + ], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)), (copyWorkgroupSize)), 65535)", + "z": 1 } }, { "id": "split_attention", "name": "GroupQueryAttention.QuantInt8DecodeSplitK", - "source": { "shader": "attn-flash-decode-splitk.wgsl.jinja", "inputs": { "layout": "\"bhsd\"" } }, - "bindings": "quantInt8SplitAttention", + "shader": "attn-flash-decode-splitk.wgsl.jinja", + "derive": { "layout": "\"bhsd\"" }, + "bindings": [ + { "arg": "queryT", "name": "query", "elementType": "vec4" }, + { "arg": "presentKeyT", "name": "key", "buffer": "read-only-storage", "elementType": "vec4" }, + { "arg": "presentValueT", "name": "value", "buffer": "read-only-storage", "elementType": "vec4" }, + "k_scale", + "v_scale", + "seqlens_k", + "partial_out", + "partial_stats", + { + "name": "params", + "struct": [ + { "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" } + ] + } + ], "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", + "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", + "derive": { "layout": "\"bhsd\"" }, + "bindings": ["partial_out_2", "partial_stats_2", "output_5"], "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": { + "derive": { "combineSubgroups": "\"_nosg\" == \"\"", "scalar": "qkvScalar", "inputVec4": "qkvInputVec4", @@ -4052,13 +2818,12 @@ "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", + "headDim": "dim(shapes.queryT, 2) / attrs.num_heads", + "headDimV4": "(dim(shapes.queryT, 2) / attrs.num_heads) / 4", + "qHiddenV4": "dim(shapes.queryT, 2) / 4", + "kvHidden": "dim(shapes.keyT, 2)", + "kvHiddenV4": "dim(shapes.keyT, 2) / 4", "numSplits": "qkvNumSplits", - "copyWorkgroupSize": "copyWorkgroupSize", "presentVec4": true, "usesF16": "qkvUsesF16", "presentElem": "qkvCacheVec4", @@ -4069,47 +2834,68 @@ { "id": "partialOut", "dtype": "float32", - "shape": "[dim(shapes.query, 0) * attrs.num_heads * qkvNumSplits * (dim(shapes.query, 2) / attrs.num_heads)]" + "shape": "[dim(shapes.queryT, 0) * attrs.num_heads * qkvNumSplits * (dim(shapes.queryT, 2) / attrs.num_heads)]" }, { "id": "partialStats", "dtype": "float32", - "shape": "[2 * dim(shapes.query, 0) * attrs.num_heads * qkvNumSplits]" + "shape": "[2 * dim(shapes.queryT, 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)" } + "shader": "attn-flash-decode-splitk.wgsl.jinja", + "derive": { "usesF16": "tensorDtypes.queryT == \"float16\"", "layout": "\"bsh\"" }, + "bindings": [ + "query_8", + { "arg": "keyT", "name": "key", "elementType": "$inputVec4" }, + { "arg": "valueT", "name": "value", "elementType": "$inputVec4" }, + "partial_out", + "partial_stats", + { + "name": "params", + "struct": [ + { "name": "kvSeq", "type": "u32", "value": "dim(shapes.keyT, 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" + } + ] + } + ], + "dispatch": { "x": "qkvNumSplits", "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": "\"bsh\"" } }, - "constants": { "usesF16": "tensorDtypes.query == \"float16\"" }, - "bindings": "qkvSplitMergeNoParams", - "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.query, 0)" } + "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", + "derive": { "usesF16": "tensorDtypes.queryT == \"float16\"", "layout": "\"bsh\"" }, + "bindings": ["partial_out_2", "partial_stats_2", "output_5"], + "dispatch": { "x": 1, "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } }, { "id": "present", "name": "GroupQueryAttention.Present", "shader": "gqa-present.wgsl.jinja", - "bindings": "qkvPresent", - "dispatch": { "threads": "numel(shapes.presentKeyT) / 4", "workgroupSize": "constants.copyWorkgroupSize" } + "bindings": ["key_2", "value_2", "present_key", "present_value", "params"], + "dispatch": { + "x": "min(ceilDiv((numel(shapes.presentKeyT) / 4), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((numel(shapes.presentKeyT) / 4), (copyWorkgroupSize)), 65535)", + "z": 1 + } } ] }, { "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, + "when": ["flashShapeOk", "qkvClusterStorageOk if \"_nosg\" == \"\" else qkvNoSgClusterStorageOk", "qkvPrefillClusterWorkgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup", "qkvPrefillClusterWorkgroupSize <= device.limits.maxComputeWorkgroupSizeX", "headDimPlan % (4 * qkvPrefillLanesPerQuery) == 0", "qkvF32ClusterRegisterGeometry if \"_nosg\" == \"\" else true", "dim(shapes.queryT, 1) >= tunables.QKV_PREFILL_MIN_QUERY_TOKENS", "subgroupQkvPrefillCluster if \"_nosg\" == \"\" else true"], "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": { + "derive": { "hasBias": false, "hasCausal": false, "headsFromParams": false, @@ -4124,12 +2910,11 @@ "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", + "headDim": "dim(shapes.queryT, 2) / attrs.num_heads", + "headDimV4": "(dim(shapes.queryT, 2) / attrs.num_heads) / 4", + "qHiddenV4": "dim(shapes.queryT, 2) / 4", + "kvHidden": "dim(shapes.keyT, 2)", + "kvHiddenV4": "dim(shapes.keyT, 2) / 4", "presentVec4": true, "TILE_Q": "qkvPrefillQueryTile", "TILE_K": "qkvNoSgTileK", @@ -4143,21 +2928,25 @@ { "id": "attention", "name": "GroupQueryAttention.Flash", - "source": { "shader": "attn-flash-prefill-cluster.wgsl.jinja", "inputs": { "layout": "\"bsh\"" } }, - "constants": { "usesF16": "tensorDtypes.query == \"float16\"" }, - "bindings": "qkvAttention", + "shader": "attn-flash-prefill-cluster.wgsl.jinja", + "derive": { "usesF16": "tensorDtypes.queryT == \"float16\"", "layout": "\"bsh\"" }, + "bindings": ["query_6", "key_2", "value_2", "output_4", "params_19"], "dispatch": { - "x": "ceilDiv(dim(shapes.query, 1), constants.TILE_Q)", + "x": "ceilDiv(dim(shapes.queryT, 1), TILE_Q)", "y": "attrs.num_heads", - "z": "dim(shapes.query, 0)" + "z": "dim(shapes.queryT, 0)" } }, { "id": "present", "name": "GroupQueryAttention.Present", "shader": "gqa-present.wgsl.jinja", - "bindings": "qkvPresent", - "dispatch": { "threads": "numel(shapes.presentKeyT) / 4", "workgroupSize": "constants.copyWorkgroupSize" } + "bindings": ["key_2", "value_2", "present_key", "present_value", "params"], + "dispatch": { + "x": "min(ceilDiv((numel(shapes.presentKeyT) / 4), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((numel(shapes.presentKeyT) / 4), (copyWorkgroupSize)), 65535)", + "z": 1 + } } ] }, @@ -4165,7 +2954,7 @@ "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": { + "derive": { "scalar": "gqaScalar", "usesF16": "tensorDtypes.queryT == \"float16\"", "hasCausal": true, @@ -4196,23 +2985,22 @@ "id": "present", "name": "GroupQueryAttention.Present", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentCopy", + "bindings": ["src_k", "src_v", "present_key_2", "present_value_2", "params_5"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "z": 1 }, "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", + "shader": "attn-flash-prefill-cluster.wgsl.jinja", + "derive": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" }, + "bindings": ["query_8", "key_7", "value_7", "attn_mask", "output_5", "params_25"], "dispatch": { - "x": "ceilDiv(dim(shapes.queryT, 1), constants.TILE_Q)", + "x": "ceilDiv(dim(shapes.queryT, 1), TILE_Q)", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } @@ -4223,7 +3011,7 @@ "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": { + "derive": { "scalar": "gqaScalar", "usesF16": "tensorDtypes.queryT == \"float16\"", "hasCausal": true, @@ -4268,10 +3056,11 @@ "id": "present", "name": "GroupQueryAttention.Present", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentCopy", + "bindings": ["src_k", "src_v", "present_key_2", "present_value_2", "params_5"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "z": 1 }, "viewAlias": [{ "input": "src_k", "output": "present_key" }, { "input": "src_v", "output": "present_value" }] }, @@ -4279,22 +3068,27 @@ "id": "qprep", "name": "GroupQueryAttention.QPrep", "shader": "gqa-qprep.wgsl.jinja", - "bindings": "qnormRotaryQprep", + "bindings": ["query_3", "cos_cache", "sin_cache", "q_norm_weight", "qout", "params_10"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)), (copyWorkgroupSize)), 65535)", + "z": 1 } }, { "id": "attention", "name": "GroupQueryAttention.FlashPrefill", - "source": { - "shader": "attn-flash-prefill-cluster.wgsl.jinja", - "inputs": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" } - }, - "bindings": "rotaryFlashAttention", + "shader": "attn-flash-prefill-cluster.wgsl.jinja", + "derive": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" }, + "bindings": [ + { "scratch": "qPrep", "name": "query", "buffer": "read-only-storage", "elementType": "vec4" }, + "key_7", + "value_7", + "output_5", + "params_26" + ], "dispatch": { - "x": "ceilDiv(dim(shapes.queryT, 1), constants.TILE_Q)", + "x": "ceilDiv(dim(shapes.queryT, 1), TILE_Q)", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } @@ -4303,10 +3097,9 @@ }, { "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": { + "derive": { "cacheSeqlens": true, "scalar": "gqaScalar", "usesF16": "tensorDtypes.queryT == \"float16\"", @@ -4342,35 +3135,38 @@ "id": "present", "name": "GroupQueryAttention.Present", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentCopy", + "bindings": ["src_k", "src_v", "present_key_2", "present_value_2", "params_5"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "z": 1 }, "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)" } + "shader": "attn-flash-decode-splitk.wgsl.jinja", + "derive": { "layout": "\"bhsd\"" }, + "bindings": ["query_8", "key_7", "value_7", "partial_out", "partial_stats", "seqlens_k", "params_27"], + "dispatch": { "x": "cachedNumSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }, + "subgroupCollectivesWidth": "portable" }, { "id": "merge", "name": "GroupQueryAttention.DecodeSplitKMerge", - "source": { "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", "inputs": { "layout": "\"bhsd\"" } }, - "bindings": "splitMergeNoParams", + "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", + "derive": { "layout": "\"bhsd\"" }, + "bindings": ["partial_out_2", "partial_stats_2", "output_5"], "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": { + "derive": { "scalar": "gqaScalar", "usesF16": "tensorDtypes.queryT == \"float16\"", "qNumHeads": "attrs.num_heads", @@ -4405,34 +3201,37 @@ "id": "present", "name": "GroupQueryAttention.Merge", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentMerge", + "bindings": ["past_k", "past_v", "new_k", "new_v", "present_key_2", "present_value_2", "params_11"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)), (copyWorkgroupSize)), 65535)", + "z": 1 } }, { "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)" } + "shader": "attn-flash-decode-splitk.wgsl.jinja", + "derive": { "layout": "\"bhsd\"" }, + "bindings": ["query_8", "key_7", "value_7", "partial_out", "partial_stats", "params_27"], + "dispatch": { "x": "cachedNumSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }, + "subgroupCollectivesWidth": "portable" }, { "id": "merge", "name": "GroupQueryAttention.DecodeSplitKMerge", - "source": { "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", "inputs": { "layout": "\"bhsd\"" } }, - "bindings": "splitMergeNoParams", + "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", + "derive": { "layout": "\"bhsd\"" }, + "bindings": ["partial_out_2", "partial_stats_2", "output_5"], "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": { + "derive": { "cacheSeqlens": true, "scalar": "gqaScalar", "usesF16": "tensorDtypes.queryT == \"float16\"", @@ -4468,34 +3267,37 @@ "id": "present", "name": "GroupQueryAttention.WindowShift", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentMergeShare", + "bindings": ["past_k", "past_v", "new_k", "new_v", "present_key_2", "present_value_2", "seqlens_k", "params_12"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)), (copyWorkgroupSize)), 65535)", + "z": 1 } }, { "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)" } + "shader": "attn-flash-decode-splitk.wgsl.jinja", + "derive": { "layout": "\"bhsd\"" }, + "bindings": ["query_8", "key_7", "value_7", "partial_out", "partial_stats", "seqlens_k", "params_27"], + "dispatch": { "x": "cachedNumSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }, + "subgroupCollectivesWidth": "portable" }, { "id": "merge", "name": "GroupQueryAttention.DecodeSplitKMerge", - "source": { "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", "inputs": { "layout": "\"bhsd\"" } }, - "bindings": "splitMergeNoParams", + "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", + "derive": { "layout": "\"bhsd\"" }, + "bindings": ["partial_out_2", "partial_stats_2", "output_5"], "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": { + "derive": { "cacheSeqlens": true, "scalar": "gqaScalar", "usesF16": "tensorDtypes.queryT == \"float16\"", @@ -4531,35 +3333,37 @@ "id": "present", "name": "GroupQueryAttention.Present", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentCopy", + "bindings": ["src_k", "src_v", "present_key_2", "present_value_2", "params_5"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "z": 1 }, "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", + "shader": "attn-flash-decode-splitk.wgsl.jinja", + "derive": { "layout": "\"bhsd\"" }, + "bindings": ["query_8", "key_7", "value_7", "partial_out", "partial_stats", "seqlens_k", "params_27"], "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", + "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", + "derive": { "layout": "\"bhsd\"" }, + "bindings": ["partial_out_2", "partial_stats_2", "output_5"], "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": { + "derive": { "scalar": "gqaScalar", "usesF16": "tensorDtypes.queryT == \"float16\"", "qNumHeads": "attrs.num_heads", @@ -4594,34 +3398,36 @@ "id": "present", "name": "GroupQueryAttention.Merge", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentMerge", + "bindings": ["past_k", "past_v", "new_k", "new_v", "present_key_2", "present_value_2", "params_11"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)), (copyWorkgroupSize)), 65535)", + "z": 1 } }, { "id": "split_attention", "name": "GroupQueryAttention.DecodeSplitK", - "source": { "shader": "attn-flash-decode-splitk.wgsl.jinja", "inputs": { "layout": "\"bhsd\"" } }, - "bindings": "splitAttention", + "shader": "attn-flash-decode-splitk.wgsl.jinja", + "derive": { "layout": "\"bhsd\"" }, + "bindings": ["query_8", "key_7", "value_7", "partial_out", "partial_stats", "params_27"], "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", + "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", + "derive": { "layout": "\"bhsd\"" }, + "bindings": ["partial_out_2", "partial_stats_2", "output_5"], "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": { + "derive": { "cacheSeqlens": true, "scalar": "gqaScalar", "usesF16": "tensorDtypes.queryT == \"float16\"", @@ -4657,34 +3463,36 @@ "id": "present", "name": "GroupQueryAttention.WindowShift", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentMergeShare", + "bindings": ["past_k", "past_v", "new_k", "new_v", "present_key_2", "present_value_2", "seqlens_k", "params_12"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)), (copyWorkgroupSize)), 65535)", + "z": 1 } }, { "id": "split_attention", "name": "GroupQueryAttention.DecodeSplitK", - "source": { "shader": "attn-flash-decode-splitk.wgsl.jinja", "inputs": { "layout": "\"bhsd\"" } }, - "bindings": "splitSeqlensAttention", + "shader": "attn-flash-decode-splitk.wgsl.jinja", + "derive": { "layout": "\"bhsd\"" }, + "bindings": ["query_8", "key_7", "value_7", "partial_out", "partial_stats", "seqlens_k", "params_27"], "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", + "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", + "derive": { "layout": "\"bhsd\"" }, + "bindings": ["partial_out_2", "partial_stats_2", "output_5"], "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": { + "derive": { "scalar": "gqaScalar", "usesF16": "tensorDtypes.queryT == \"float16\"", "hasCausal": true, @@ -4714,35 +3522,34 @@ "id": "present", "name": "GroupQueryAttention.Present", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentCopy", + "bindings": ["src_k", "src_v", "present_key_2", "present_value_2", "params_5"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "z": 1 }, "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", + "shader": "attn-flash-prefill-cluster.wgsl.jinja", + "derive": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" }, + "bindings": ["query_8", "key_7", "value_7", "output_5", "params_26"], "dispatch": { - "x": "ceilDiv(dim(shapes.queryT, 1), constants.TILE_Q)", + "x": "ceilDiv(dim(shapes.queryT, 1), TILE_Q)", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" - } + }, + "subgroupCollectivesWidth": "portable" } ] }, { "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": { + "derive": { "scalar": "gqaScalar", "usesF16": "tensorDtypes.queryT == \"float16\"", "hasCausal": true, @@ -4772,34 +3579,33 @@ "id": "present", "name": "GroupQueryAttention.Merge", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentMerge", + "bindings": ["past_k", "past_v", "new_k", "new_v", "present_key_2", "present_value_2", "params_11"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)), (copyWorkgroupSize)), 65535)", + "z": 1 } }, { "id": "attention", "name": "GroupQueryAttention.FlashPrefill", - "source": { - "shader": "attn-flash-prefill-cluster.wgsl.jinja", - "inputs": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" } - }, - "bindings": "flashPrefillNew", + "shader": "attn-flash-prefill-cluster.wgsl.jinja", + "derive": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" }, + "bindings": ["query_8", "key_7", "value_7", "output_5", "params_28"], "dispatch": { - "x": "ceilDiv(dim(shapes.queryT, 1), constants.TILE_Q)", + "x": "ceilDiv(dim(shapes.queryT, 1), TILE_Q)", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" - } + }, + "subgroupCollectivesWidth": "portable" } ] }, { "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": { + "derive": { "scalar": "gqaScalar", "usesF16": "tensorDtypes.queryT == \"float16\"", "hasCausal": true, @@ -4835,35 +3641,34 @@ "id": "present", "name": "GroupQueryAttention.Present", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentCopy", + "bindings": ["src_k", "src_v", "present_key_2", "present_value_2", "params_5"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "z": 1 }, "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", + "shader": "attn-flash-prefill-cluster.wgsl.jinja", + "derive": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" }, + "bindings": ["query_8", "key_7", "value_7", "cos_cache", "sin_cache", "output_5", "params_26"], "dispatch": { - "x": "ceilDiv(dim(shapes.queryT, 1), constants.TILE_Q)", + "x": "ceilDiv(dim(shapes.queryT, 1), TILE_Q)", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" - } + }, + "subgroupCollectivesWidth": "portable" } ] }, { "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": { + "derive": { "scalar": "gqaScalar", "usesF16": "tensorDtypes.queryT == \"float16\"", "hasCausal": true, @@ -4894,35 +3699,34 @@ "id": "present", "name": "GroupQueryAttention.Present", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentCopy", + "bindings": ["src_k", "src_v", "present_key_2", "present_value_2", "params_5"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "z": 1 }, "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", + "shader": "attn-flash-prefill-cluster.wgsl.jinja", + "derive": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" }, + "bindings": ["query_8", "key_7", "value_7", "output_5", "params_29"], "dispatch": { - "x": "ceilDiv(dim(shapes.queryT, 1), constants.TILE_Q)", + "x": "ceilDiv(dim(shapes.queryT, 1), TILE_Q)", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" - } + }, + "subgroupCollectivesWidth": "portable" } ] }, { "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": { + "derive": { "scalar": "gqaScalar", "usesF16": "tensorDtypes.queryT == \"float16\"", "hasCausal": true, @@ -4953,35 +3757,34 @@ "id": "present", "name": "GroupQueryAttention.Present", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentCopy", + "bindings": ["src_k", "src_v", "present_key_2", "present_value_2", "params_5"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "z": 1 }, "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", + "shader": "attn-flash-prefill-cluster.wgsl.jinja", + "derive": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" }, + "bindings": ["query_8", "key_7", "value_7", "head_sink", "output_5", "params_26"], "dispatch": { - "x": "ceilDiv(dim(shapes.queryT, 1), constants.TILE_Q)", + "x": "ceilDiv(dim(shapes.queryT, 1), TILE_Q)", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" - } + }, + "subgroupCollectivesWidth": "portable" } ] }, { "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": { + "derive": { "scalar": "gqaScalar", "usesF16": "tensorDtypes.queryT == \"float16\"", "hasCausal": true, @@ -5012,35 +3815,34 @@ "id": "present", "name": "GroupQueryAttention.Present", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentCopy", + "bindings": ["src_k", "src_v", "present_key_2", "present_value_2", "params_5"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "z": 1 }, "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", + "shader": "attn-flash-prefill-cluster.wgsl.jinja", + "derive": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" }, + "bindings": ["query_8", "key_7", "value_7", "attn_mask", "head_sink", "output_5", "params_25"], "dispatch": { - "x": "ceilDiv(dim(shapes.queryT, 1), constants.TILE_Q)", + "x": "ceilDiv(dim(shapes.queryT, 1), TILE_Q)", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" - } + }, + "subgroupCollectivesWidth": "portable" } ] }, { "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": { + "derive": { "scalar": "gqaScalar", "usesF16": "tensorDtypes.queryT == \"float16\"", "hasCausal": true, @@ -5072,34 +3874,33 @@ "id": "present", "name": "GroupQueryAttention.WindowShift", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentMergeShare", + "bindings": ["past_k", "past_v", "new_k", "new_v", "present_key_2", "present_value_2", "seqlens_k", "params_12"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)), (copyWorkgroupSize)), 65535)", + "z": 1 } }, { "id": "attention", "name": "GroupQueryAttention.FlashPrefill", - "source": { - "shader": "attn-flash-prefill-cluster.wgsl.jinja", - "inputs": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" } - }, - "bindings": "flashPrefillSeqlens", + "shader": "attn-flash-prefill-cluster.wgsl.jinja", + "derive": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" }, + "bindings": ["query_8", "key_7", "value_7", "output_5", "seqlens_k", "params_28"], "dispatch": { - "x": "ceilDiv(dim(shapes.queryT, 1), constants.TILE_Q)", + "x": "ceilDiv(dim(shapes.queryT, 1), TILE_Q)", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" - } + }, + "subgroupCollectivesWidth": "portable" } ] }, { "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": { + "derive": { "scalar": "gqaScalar", "usesF16": "tensorDtypes.queryT == \"float16\"", "hasCausal": true, @@ -5129,23 +3930,22 @@ "id": "present", "name": "GroupQueryAttention.Present", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentCopy", + "bindings": ["src_k", "src_v", "present_key_2", "present_value_2", "params_5"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "z": 1 }, "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", + "shader": "attn-flash-prefill-cluster.wgsl.jinja", + "derive": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" }, + "bindings": ["query_8", "key_7", "value_7", "output_5", "params_26"], "dispatch": { - "x": "ceilDiv(dim(shapes.queryT, 1), constants.TILE_Q)", + "x": "ceilDiv(dim(shapes.queryT, 1), TILE_Q)", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } @@ -5154,10 +3954,9 @@ }, { "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": { + "derive": { "scalar": "gqaScalar", "usesF16": "tensorDtypes.queryT == \"float16\"", "hasCausal": true, @@ -5187,22 +3986,21 @@ "id": "present", "name": "GroupQueryAttention.Merge", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentMerge", + "bindings": ["past_k", "past_v", "new_k", "new_v", "present_key_2", "present_value_2", "params_11"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)), (copyWorkgroupSize)), 65535)", + "z": 1 } }, { "id": "attention", "name": "GroupQueryAttention.FlashPrefill", - "source": { - "shader": "attn-flash-prefill-cluster.wgsl.jinja", - "inputs": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" } - }, - "bindings": "flashPrefillNew", + "shader": "attn-flash-prefill-cluster.wgsl.jinja", + "derive": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" }, + "bindings": ["query_8", "key_7", "value_7", "output_5", "params_28"], "dispatch": { - "x": "ceilDiv(dim(shapes.queryT, 1), constants.TILE_Q)", + "x": "ceilDiv(dim(shapes.queryT, 1), TILE_Q)", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } @@ -5211,10 +4009,9 @@ }, { "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": { + "derive": { "scalar": "gqaScalar", "usesF16": "tensorDtypes.queryT == \"float16\"", "hasCausal": true, @@ -5250,23 +4047,22 @@ "id": "present", "name": "GroupQueryAttention.Present", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentCopy", + "bindings": ["src_k", "src_v", "present_key_2", "present_value_2", "params_5"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "z": 1 }, "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", + "shader": "attn-flash-prefill-cluster.wgsl.jinja", + "derive": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" }, + "bindings": ["query_8", "key_7", "value_7", "cos_cache", "sin_cache", "output_5", "params_26"], "dispatch": { - "x": "ceilDiv(dim(shapes.queryT, 1), constants.TILE_Q)", + "x": "ceilDiv(dim(shapes.queryT, 1), TILE_Q)", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } @@ -5275,10 +4071,9 @@ }, { "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": { + "derive": { "scalar": "gqaScalar", "usesF16": "tensorDtypes.queryT == \"float16\"", "hasCausal": true, @@ -5309,23 +4104,22 @@ "id": "present", "name": "GroupQueryAttention.Present", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentCopy", + "bindings": ["src_k", "src_v", "present_key_2", "present_value_2", "params_5"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "z": 1 }, "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", + "shader": "attn-flash-prefill-cluster.wgsl.jinja", + "derive": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" }, + "bindings": ["query_8", "key_7", "value_7", "output_5", "params_29"], "dispatch": { - "x": "ceilDiv(dim(shapes.queryT, 1), constants.TILE_Q)", + "x": "ceilDiv(dim(shapes.queryT, 1), TILE_Q)", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } @@ -5334,10 +4128,9 @@ }, { "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": { + "derive": { "scalar": "gqaScalar", "usesF16": "tensorDtypes.queryT == \"float16\"", "hasCausal": true, @@ -5368,23 +4161,22 @@ "id": "present", "name": "GroupQueryAttention.Present", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentCopy", + "bindings": ["src_k", "src_v", "present_key_2", "present_value_2", "params_5"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "z": 1 }, "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", + "shader": "attn-flash-prefill-cluster.wgsl.jinja", + "derive": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" }, + "bindings": ["query_8", "key_7", "value_7", "head_sink", "output_5", "params_26"], "dispatch": { - "x": "ceilDiv(dim(shapes.queryT, 1), constants.TILE_Q)", + "x": "ceilDiv(dim(shapes.queryT, 1), TILE_Q)", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } @@ -5393,10 +4185,9 @@ }, { "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": { + "derive": { "scalar": "gqaScalar", "usesF16": "tensorDtypes.queryT == \"float16\"", "hasCausal": true, @@ -5427,23 +4218,22 @@ "id": "present", "name": "GroupQueryAttention.Present", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentCopy", + "bindings": ["src_k", "src_v", "present_key_2", "present_value_2", "params_5"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "z": 1 }, "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", + "shader": "attn-flash-prefill-cluster.wgsl.jinja", + "derive": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" }, + "bindings": ["query_8", "key_7", "value_7", "attn_mask", "head_sink", "output_5", "params_25"], "dispatch": { - "x": "ceilDiv(dim(shapes.queryT, 1), constants.TILE_Q)", + "x": "ceilDiv(dim(shapes.queryT, 1), TILE_Q)", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } @@ -5452,10 +4242,9 @@ }, { "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": { + "derive": { "scalar": "gqaScalar", "usesF16": "tensorDtypes.queryT == \"float16\"", "hasCausal": true, @@ -5487,22 +4276,21 @@ "id": "present", "name": "GroupQueryAttention.WindowShift", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentMergeShare", + "bindings": ["past_k", "past_v", "new_k", "new_v", "present_key_2", "present_value_2", "seqlens_k", "params_12"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)), (copyWorkgroupSize)), 65535)", + "z": 1 } }, { "id": "attention", "name": "GroupQueryAttention.FlashPrefill", - "source": { - "shader": "attn-flash-prefill-cluster.wgsl.jinja", - "inputs": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" } - }, - "bindings": "flashPrefillSeqlens", + "shader": "attn-flash-prefill-cluster.wgsl.jinja", + "derive": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" }, + "bindings": ["query_8", "key_7", "value_7", "output_5", "seqlens_k", "params_28"], "dispatch": { - "x": "ceilDiv(dim(shapes.queryT, 1), constants.TILE_Q)", + "x": "ceilDiv(dim(shapes.queryT, 1), TILE_Q)", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" } @@ -5512,9 +4300,9 @@ { "id": "qkv_present_flash_q32_broadcast", "priority": 30, + "when": ["subgroupsWave32", "device.features.has(\"shader-f16\")", "tensorDtypes.queryT == \"float16\"", "flashShapeOk", "headDimPlan % 32 == 0", "headDimPlan >= 64", "headDimPlan <= 256", "dim(shapes.queryT, 1) >= tunables.QKV_PREFILL_MIN_QUERY_TOKENS", "attrs.local_window_size <= 0"], "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": { + "derive": { "hasBias": false, "hasCausal": false, "usesF16": true, @@ -5527,11 +4315,10 @@ "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", + "headDim": "dim(shapes.queryT, 2) / attrs.num_heads", + "headDimV4": "(dim(shapes.queryT, 2) / attrs.num_heads) / 4", + "kvHidden": "dim(shapes.keyT, 2)", + "kvHiddenV4": "dim(shapes.keyT, 2) / 4", "presentVec4": true, "presentElem": "qkvCacheVec4", "kStep": 64, @@ -5541,26 +4328,48 @@ { "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)" } + "shader": "attn-flash-q32-broadcast.wgsl.jinja", + "derive": { "usesF16": "tensorDtypes.queryT == \"float16\"", "layout": "\"bsh\"" }, + "bindings": [ + "query_6", + "key_2", + "value_2", + "output_4", + { + "name": "params", + "struct": [ + { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, + { "name": "kvSeq", "type": "u32", "value": "dim(shapes.keyT, 1)" }, + { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" } + ] + } + ], + "dispatch": { + "x": "ceilDiv(dim(shapes.queryT, 1), 32)", + "y": "attrs.num_heads", + "z": "dim(shapes.queryT, 0)" + }, + "subgroupCollectivesWidth": 32 }, { "id": "present", "name": "GroupQueryAttention.Present", "shader": "gqa-present.wgsl.jinja", - "bindings": "qkvPresent", - "dispatch": { "threads": "numel(shapes.presentKeyT) / 4", "workgroupSize": "constants.copyWorkgroupSize" } + "bindings": ["key_2", "value_2", "present_key", "present_value", "params"], + "dispatch": { + "x": "min(ceilDiv((numel(shapes.presentKeyT) / 4), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((numel(shapes.presentKeyT) / 4), (copyWorkgroupSize)), 65535)", + "z": 1 + } } ] }, { "id": "qkv_present_flash_q32_shared", "priority": 29, + "when": ["tensorDtypes.queryT == \"float16\"", "flashShapeOk", "headDimPlan % 4 == 0", "headDimPlan >= 64", "headDimPlan <= 256", "dim(shapes.queryT, 1) >= tunables.QKV_PREFILL_MIN_QUERY_TOKENS", "attrs.local_window_size <= 0", "ceilDiv(dim(shapes.queryT, 1), 64) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "(headDimPlan / 4) * 32 * 16 <= device.limits.maxComputeWorkgroupStorageSize"], "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": { + "derive": { "hasBias": false, "hasCausal": false, "usesF16": true, @@ -5573,11 +4382,10 @@ "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", + "headDim": "dim(shapes.queryT, 2) / attrs.num_heads", + "headDimV4": "(dim(shapes.queryT, 2) / attrs.num_heads) / 4", + "kvHidden": "dim(shapes.keyT, 2)", + "kvHiddenV4": "dim(shapes.keyT, 2) / 4", "presentVec4": true, "presentElem": "qkvCacheVec4", "kStep": 32, @@ -5588,20 +4396,38 @@ { "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)" } + "shader": "attn-flash-q32-broadcast.wgsl.jinja", + "derive": { "usesF16": "tensorDtypes.queryT == \"float16\"", "layout": "\"bsh\"", "useSubgroups": "false" }, + "bindings": [ + "query_6", + "key_2", + "value_2", + "output_4", + { + "name": "params", + "struct": [ + { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" }, + { "name": "kvSeq", "type": "u32", "value": "dim(shapes.keyT, 1)" }, + { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" } + ] + } + ], + "dispatch": { + "x": "ceilDiv(dim(shapes.queryT, 1), 64)", + "y": "attrs.num_heads", + "z": "dim(shapes.queryT, 0)" + } }, { "id": "present", "name": "GroupQueryAttention.Present", "shader": "gqa-present.wgsl.jinja", - "bindings": "qkvPresent", - "dispatch": { "threads": "numel(shapes.presentKeyT) / 4", "workgroupSize": "constants.copyWorkgroupSize" } + "bindings": ["key_2", "value_2", "present_key", "present_value", "params"], + "dispatch": { + "x": "min(ceilDiv((numel(shapes.presentKeyT) / 4), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((numel(shapes.presentKeyT) / 4), (copyWorkgroupSize)), 65535)", + "z": 1 + } } ] }, @@ -5609,7 +4435,7 @@ "id": "past_kv", "priority": 30, "when": ["sharedKvFloatOk", "gqaScalarDispatchOk", "noAuxAttentionInputs", "gqaHeadDim % 2 == 0"], - "constants": { + "derive": { "useSeqlens": true, "headDim": "gqaHeadDim", "cooperative": "gqaCoop", @@ -5627,17 +4453,18 @@ "id": "attention", "name": "GroupQueryAttention.Attn", "shader": "gqa-attention.wgsl.jinja", - "bindings": "cachedAttention", - "dispatch": { "workgroups": "constants.dispatchUnits" } + "bindings": ["query_3", "kcache_3", "vcache_3", "output_3", "seqlens_k", "params_31"], + "dispatch": { "x": "min(dispatchUnits, 65535)", "y": "ceilDiv(dispatchUnits, 65535)", "z": 1 } }, { "id": "present", "name": "GroupQueryAttention.Present", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentCopy", + "bindings": ["src_k", "src_v", "present_key_2", "present_value_2", "params_5"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "z": 1 }, "viewAlias": [{ "input": "src_k", "output": "present_key" }, { "input": "src_v", "output": "present_value" }] } @@ -5647,7 +4474,7 @@ "id": "past_kv_rotary", "priority": 31, "when": ["sharedKvFloatOk", "gqaScalarDispatchOk", "gqaHeadDim % 2 == 0", "present.cosCacheT and rotaryRequested", "not present.qNormWeightT"], - "constants": { + "derive": { "useSeqlens": true, "headDim": "gqaHeadDim", "cooperative": "gqaCoop", @@ -5668,17 +4495,18 @@ "id": "attention", "name": "GroupQueryAttention.Attn", "shader": "gqa-attention.wgsl.jinja", - "bindings": "cachedRotaryAttention", - "dispatch": { "workgroups": "constants.dispatchUnits" } + "bindings": ["query_3", "kcache_3", "vcache_3", "cos_cache", "sin_cache", "output_3", "seqlens_k", "params_31"], + "dispatch": { "x": "min(dispatchUnits, 65535)", "y": "ceilDiv(dispatchUnits, 65535)", "z": 1 } }, { "id": "present", "name": "GroupQueryAttention.Present", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentCopy", + "bindings": ["src_k", "src_v", "present_key_2", "present_value_2", "params_5"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "z": 1 }, "viewAlias": [{ "input": "src_k", "output": "present_key" }, { "input": "src_v", "output": "present_value" }] } @@ -5688,7 +4516,7 @@ "id": "past_kv_qnorm_rotary", "priority": 32, "when": ["sharedKvFloatOk", "gqaScalarDispatchOk", "gqaHeadDim % 2 == 0", "present.cosCacheT and rotaryRequested", "present.qNormWeightT"], - "constants": { + "derive": { "useSeqlens": true, "headDim": "gqaHeadDim", "cooperative": "gqaCoop", @@ -5712,17 +4540,18 @@ "id": "attention", "name": "GroupQueryAttention.Attn", "shader": "gqa-attention.wgsl.jinja", - "bindings": "cachedQnormRotaryAttention", - "dispatch": { "workgroups": "constants.dispatchUnits" } + "bindings": ["query_3", "kcache_3", "vcache_3", "cos_cache", "sin_cache", "q_norm_weight", "output_3", "seqlens_k", "params_31"], + "dispatch": { "x": "min(dispatchUnits, 65535)", "y": "ceilDiv(dispatchUnits, 65535)", "z": 1 } }, { "id": "present", "name": "GroupQueryAttention.Present", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentCopy", + "bindings": ["src_k", "src_v", "present_key_2", "present_value_2", "params_5"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "z": 1 }, "viewAlias": [{ "input": "src_k", "output": "present_key" }, { "input": "src_v", "output": "present_value" }] } @@ -5732,7 +4561,7 @@ "id": "new_kv_past", "priority": 30, "when": ["newKvPastOk", "gqaScalarDispatchOk", "noAuxAttentionInputs", "gqaHeadDim % 2 == 0"], - "constants": { + "derive": { "headDim": "gqaHeadDim", "qHeads": "attrs.num_heads", "kvHeads": "attrs.kv_num_heads", @@ -5749,18 +4578,19 @@ "id": "present", "name": "GroupQueryAttention.Merge", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentMerge", + "bindings": ["past_k", "past_v", "new_k", "new_v", "present_key_2", "present_value_2", "params_11"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)), (copyWorkgroupSize)), 65535)", + "z": 1 } }, { "id": "attention", "name": "GroupQueryAttention.Attn", "shader": "gqa-attention.wgsl.jinja", - "bindings": "mergedAttention", - "dispatch": { "workgroups": "constants.dispatchUnits" } + "bindings": ["query_3", "kcache", "vcache", "output_3", "params_16"], + "dispatch": { "x": "min(dispatchUnits, 65535)", "y": "ceilDiv(dispatchUnits, 65535)", "z": 1 } } ] }, @@ -5768,7 +4598,7 @@ "id": "window_shift_append", "priority": 30, "when": ["windowShiftOk", "gqaScalarDispatchOk", "noAuxAttentionInputs", "gqaHeadDim % 2 == 0"], - "constants": { + "derive": { "useSeqlens": true, "headDim": "gqaHeadDim", "cooperative": "gqaCoop", @@ -5787,18 +4617,19 @@ "id": "present", "name": "GroupQueryAttention.WindowShift", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentMergeShare", + "bindings": ["past_k", "past_v", "new_k", "new_v", "present_key_2", "present_value_2", "seqlens_k", "params_12"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)), (copyWorkgroupSize)), 65535)", + "z": 1 } }, { "id": "attention", "name": "GroupQueryAttention.Attn", "shader": "gqa-attention.wgsl.jinja", - "bindings": "mergedSeqlensAttention", - "dispatch": { "workgroups": "constants.dispatchUnits" } + "bindings": ["query_3", "kcache", "vcache", "output_3", "seqlens_k", "params_16"], + "dispatch": { "x": "min(dispatchUnits, 65535)", "y": "ceilDiv(dispatchUnits, 65535)", "z": 1 } } ] }, @@ -5806,7 +4637,7 @@ "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": { + "derive": { "headDim": "gqaHeadDim", "cooperative": "gqaCoop", "dispatchUnits": "gqaDispatchUnits", @@ -5826,18 +4657,19 @@ "id": "present", "name": "GroupQueryAttention.Present", "shader": "gqa-present.wgsl.jinja", - "bindings": "newQnormRotaryPresent", + "bindings": ["src_k_2", "src_v_2", "cos_cache_2", "sin_cache_2", "k_norm_weight", "present_key_5", "present_value_5", "params_32"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)), (copyWorkgroupSize)), 65535)", + "z": 1 } }, { "id": "attention", "name": "GroupQueryAttention.Attn", "shader": "gqa-attention.wgsl.jinja", - "bindings": "newQnormRotaryAttention", - "dispatch": { "workgroups": "constants.dispatchUnits" } + "bindings": ["query", "kcache_4", "vcache_4", "cos_cache_2", "sin_cache_2", "q_norm_weight_2", "output", "params_16"], + "dispatch": { "x": "min(dispatchUnits, 65535)", "y": "ceilDiv(dispatchUnits, 65535)", "z": 1 } } ] }, @@ -5845,7 +4677,7 @@ "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": { + "derive": { "useSeqlens": true, "headDim": "gqaHeadDim", "cooperative": "gqaCoop", @@ -5863,17 +4695,18 @@ "id": "attention", "name": "GroupQueryAttention.Attn", "shader": "gqa-attention.wgsl.jinja", - "bindings": "cachedBiasAttention", - "dispatch": { "workgroups": "constants.dispatchUnits" } + "bindings": ["query", "kcache_5", "vcache_5", "attn_bias", "output", "seqlens_k", "params_33"], + "dispatch": { "x": "min(dispatchUnits, 65535)", "y": "ceilDiv(dispatchUnits, 65535)", "z": 1 } }, { "id": "present", "name": "GroupQueryAttention.Present", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentCopy", + "bindings": ["src_k", "src_v", "present_key_2", "present_value_2", "params_5"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "z": 1 }, "viewAlias": [{ "input": "src_k", "output": "present_key" }, { "input": "src_v", "output": "present_value" }] } @@ -5883,7 +4716,7 @@ "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": { + "derive": { "useSeqlens": true, "headDim": "gqaHeadDim", "cooperative": "gqaCoop", @@ -5901,17 +4734,18 @@ "id": "attention", "name": "GroupQueryAttention.Attn", "shader": "gqa-attention.wgsl.jinja", - "bindings": "cachedHeadSinkAttention", - "dispatch": { "workgroups": "constants.dispatchUnits" } + "bindings": ["query", "kcache_5", "vcache_5", "head_sink", "output", "seqlens_k", "params_34"], + "dispatch": { "x": "min(dispatchUnits, 65535)", "y": "ceilDiv(dispatchUnits, 65535)", "z": 1 } }, { "id": "present", "name": "GroupQueryAttention.Present", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentCopy", + "bindings": ["src_k", "src_v", "present_key_2", "present_value_2", "params_5"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "z": 1 }, "viewAlias": [{ "input": "src_k", "output": "present_key" }, { "input": "src_v", "output": "present_value" }] } @@ -5921,7 +4755,7 @@ "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": { + "derive": { "useSeqlens": true, "headDim": "gqaHeadDim", "cooperative": "gqaCoop", @@ -5940,17 +4774,18 @@ "id": "attention", "name": "GroupQueryAttention.Attn", "shader": "gqa-attention.wgsl.jinja", - "bindings": "cachedBiasHeadSinkAttention", - "dispatch": { "workgroups": "constants.dispatchUnits" } + "bindings": ["query", "kcache_5", "vcache_5", "attn_bias", "head_sink", "output", "seqlens_k", "params_35"], + "dispatch": { "x": "min(dispatchUnits, 65535)", "y": "ceilDiv(dispatchUnits, 65535)", "z": 1 } }, { "id": "present", "name": "GroupQueryAttention.Present", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentCopy", + "bindings": ["src_k", "src_v", "present_key_2", "present_value_2", "params_5"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)), (copyWorkgroupSize)), 65535)", + "z": 1 }, "viewAlias": [{ "input": "src_k", "output": "present_key" }, { "input": "src_v", "output": "present_value" }] } @@ -5958,10 +4793,9 @@ }, { "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": { + "derive": { "scalar": "gqaScalar", "usesF16": false, "hasCausal": true, @@ -6002,34 +4836,33 @@ "id": "present", "name": "GroupQueryAttention.QuantPresent", "shader": "gqa-present.wgsl.jinja", - "bindings": "quantPresent", + "bindings": ["src_k_2", "src_v_2", "k_scale", "v_scale", "present_key_3", "present_value_3", "params_20"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)), (copyWorkgroupSize)), 65535)", + "z": 1 } }, { "id": "attention", "name": "GroupQueryAttention.QuantFlashPrefill", - "source": { - "shader": "attn-flash-prefill-cluster.wgsl.jinja", - "inputs": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" } - }, - "bindings": "flashPrefillQuant", + "shader": "attn-flash-prefill-cluster.wgsl.jinja", + "derive": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" }, + "bindings": ["query_8", "key_8", "value_8", "k_scale", "v_scale", "output_5", "params_36"], "dispatch": { "x": "ceilDiv(dim(shapes.queryT, 1), cachedPrefillQueryTile)", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" - } + }, + "subgroupCollectivesWidth": "portable" } ] }, { "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": { + "derive": { "scalar": "gqaScalar", "usesF16": false, "hasCausal": true, @@ -6070,34 +4903,33 @@ "id": "present", "name": "GroupQueryAttention.QuantPresent", "shader": "gqa-present.wgsl.jinja", - "bindings": "quantPresent", + "bindings": ["src_k_2", "src_v_2", "k_scale", "v_scale", "present_key_3", "present_value_3", "params_20"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)), (copyWorkgroupSize)), 65535)", + "z": 1 } }, { "id": "attention", "name": "GroupQueryAttention.QuantFlashPrefill", - "source": { - "shader": "attn-flash-prefill-cluster.wgsl.jinja", - "inputs": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" } - }, - "bindings": "flashPrefillQuant", + "shader": "attn-flash-prefill-cluster.wgsl.jinja", + "derive": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" }, + "bindings": ["query_8", "key_8", "value_8", "k_scale", "v_scale", "output_5", "params_36"], "dispatch": { "x": "ceilDiv(dim(shapes.queryT, 1), cachedPrefillQueryTile)", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" - } + }, + "subgroupCollectivesWidth": "portable" } ] }, { "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": { + "derive": { "scalar": "gqaScalar", "usesF16": false, "hasCausal": true, @@ -6138,20 +4970,19 @@ "id": "present", "name": "GroupQueryAttention.QuantPresent", "shader": "gqa-present.wgsl.jinja", - "bindings": "quantPresent", + "bindings": ["src_k_2", "src_v_2", "k_scale", "v_scale", "present_key_3", "present_value_3", "params_20"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)), (copyWorkgroupSize)), 65535)", + "z": 1 } }, { "id": "attention", "name": "GroupQueryAttention.QuantFlashPrefill", - "source": { - "shader": "attn-flash-prefill-cluster.wgsl.jinja", - "inputs": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" } - }, - "bindings": "flashPrefillQuant", + "shader": "attn-flash-prefill-cluster.wgsl.jinja", + "derive": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" }, + "bindings": ["query_8", "key_8", "value_8", "k_scale", "v_scale", "output_5", "params_36"], "dispatch": { "x": "ceilDiv(dim(shapes.queryT, 1), cachedPrefillQueryTile)", "y": "attrs.num_heads", @@ -6162,10 +4993,9 @@ }, { "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": { + "derive": { "scalar": "gqaScalar", "usesF16": false, "hasCausal": true, @@ -6206,20 +5036,19 @@ "id": "present", "name": "GroupQueryAttention.QuantPresent", "shader": "gqa-present.wgsl.jinja", - "bindings": "quantPresent", + "bindings": ["src_k_2", "src_v_2", "k_scale", "v_scale", "present_key_3", "present_value_3", "params_20"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)), (copyWorkgroupSize)), 65535)", + "z": 1 } }, { "id": "attention", "name": "GroupQueryAttention.QuantFlashPrefill", - "source": { - "shader": "attn-flash-prefill-cluster.wgsl.jinja", - "inputs": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" } - }, - "bindings": "flashPrefillQuant", + "shader": "attn-flash-prefill-cluster.wgsl.jinja", + "derive": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" }, + "bindings": ["query_8", "key_8", "value_8", "k_scale", "v_scale", "output_5", "params_36"], "dispatch": { "x": "ceilDiv(dim(shapes.queryT, 1), cachedPrefillQueryTile)", "y": "attrs.num_heads", @@ -6230,10 +5059,9 @@ }, { "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": { + "derive": { "cacheSeqlens": "true", "scalar": "gqaScalar", "usesF16": "tensorDtypes.queryT == \"float16\"", @@ -6269,11 +5097,12 @@ "id": "present_retain", "name": "GroupQueryAttention.MergeShareRetain", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentMergeShareRetain", - "constants": { "shareRegion": "\"retain\"" }, + "derive": { "shareRegion": "\"retain\"" }, + "bindings": ["past_k", "past_v", "present_key_2", "present_value_2", "seqlens_k", "params_12"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)), (copyWorkgroupSize)), 65535)", + "z": 1 }, "viewAlias": [ { "input": "past_k", "output": "present_key" }, @@ -6284,35 +5113,38 @@ "id": "present_append", "name": "GroupQueryAttention.MergeShareAppend", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentMergeShareAppend", - "constants": { "shareRegion": "\"append\"" }, + "derive": { "shareRegion": "\"append\"" }, + "bindings": ["new_k", "new_v", "present_key_2", "present_value_2", "seqlens_k", "params_15"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)), (copyWorkgroupSize)), 65535)", + "z": 1 } }, { "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)" } + "shader": "attn-flash-decode-splitk.wgsl.jinja", + "derive": { "layout": "\"bhsd\"" }, + "bindings": ["query_8", "key_7", "value_7", "partial_out", "partial_stats", "seqlens_k", "params_27"], + "dispatch": { "x": "cachedNumSplits", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }, + "subgroupCollectivesWidth": "portable" }, { "id": "merge", "name": "GroupQueryAttention.DecodeSplitKMerge", - "source": { "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", "inputs": { "layout": "\"bhsd\"" } }, - "bindings": "splitMergeNoParams", + "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", + "derive": { "layout": "\"bhsd\"" }, + "bindings": ["partial_out_2", "partial_stats_2", "output_5"], "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": { + "derive": { "cacheSeqlens": "true", "scalar": "gqaScalar", "usesF16": "tensorDtypes.queryT == \"float16\"", @@ -6348,11 +5180,12 @@ "id": "present_retain", "name": "GroupQueryAttention.MergeShareRetain", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentMergeShareRetain", - "constants": { "shareRegion": "\"retain\"" }, + "derive": { "shareRegion": "\"retain\"" }, + "bindings": ["past_k", "past_v", "present_key_2", "present_value_2", "seqlens_k", "params_12"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)), (copyWorkgroupSize)), 65535)", + "z": 1 }, "viewAlias": [ { "input": "past_k", "output": "present_key" }, @@ -6363,35 +5196,37 @@ "id": "present_append", "name": "GroupQueryAttention.MergeShareAppend", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentMergeShareAppend", - "constants": { "shareRegion": "\"append\"" }, + "derive": { "shareRegion": "\"append\"" }, + "bindings": ["new_k", "new_v", "present_key_2", "present_value_2", "seqlens_k", "params_15"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)), (copyWorkgroupSize)), 65535)", + "z": 1 } }, { "id": "split_attention", "name": "GroupQueryAttention.DecodeSplitK", - "source": { "shader": "attn-flash-decode-splitk.wgsl.jinja", "inputs": { "layout": "\"bhsd\"" } }, - "bindings": "splitSeqlensAttention", + "shader": "attn-flash-decode-splitk.wgsl.jinja", + "derive": { "layout": "\"bhsd\"" }, + "bindings": ["query_8", "key_7", "value_7", "partial_out", "partial_stats", "seqlens_k", "params_27"], "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", + "shader": "attn-flash-decode-splitk-merge.wgsl.jinja", + "derive": { "layout": "\"bhsd\"" }, + "bindings": ["partial_out_2", "partial_stats_2", "output_5"], "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": { + "derive": { "scalar": "gqaScalar", "usesF16": "tensorDtypes.queryT == \"float16\"", "hasCausal": true, @@ -6422,11 +5257,12 @@ "id": "present_retain", "name": "GroupQueryAttention.MergeShareRetain", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentMergeShareRetain", - "constants": { "shareRegion": "\"retain\"" }, + "derive": { "shareRegion": "\"retain\"" }, + "bindings": ["past_k", "past_v", "present_key_2", "present_value_2", "seqlens_k", "params_12"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)), (copyWorkgroupSize)), 65535)", + "z": 1 }, "viewAlias": [ { "input": "past_k", "output": "present_key" }, @@ -6437,35 +5273,34 @@ "id": "present_append", "name": "GroupQueryAttention.MergeShareAppend", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentMergeShareAppend", - "constants": { "shareRegion": "\"append\"" }, + "derive": { "shareRegion": "\"append\"" }, + "bindings": ["new_k", "new_v", "present_key_2", "present_value_2", "seqlens_k", "params_15"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)), (copyWorkgroupSize)), 65535)", + "z": 1 } }, { "id": "attention", "name": "GroupQueryAttention.FlashPrefill", - "source": { - "shader": "attn-flash-prefill-cluster.wgsl.jinja", - "inputs": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" } - }, - "bindings": "flashPrefillSeqlens", + "shader": "attn-flash-prefill-cluster.wgsl.jinja", + "derive": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" }, + "bindings": ["query_8", "key_7", "value_7", "output_5", "seqlens_k", "params_28"], "dispatch": { - "x": "ceilDiv(dim(shapes.queryT, 1), constants.TILE_Q)", + "x": "ceilDiv(dim(shapes.queryT, 1), TILE_Q)", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" - } + }, + "subgroupCollectivesWidth": "portable" } ] }, { "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": { + "derive": { "scalar": "gqaScalar", "usesF16": "tensorDtypes.queryT == \"float16\"", "hasCausal": true, @@ -6496,11 +5331,12 @@ "id": "present_retain", "name": "GroupQueryAttention.MergeShareRetain", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentMergeShareRetain", - "constants": { "shareRegion": "\"retain\"" }, + "derive": { "shareRegion": "\"retain\"" }, + "bindings": ["past_k", "past_v", "present_key_2", "present_value_2", "seqlens_k", "params_12"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)), (copyWorkgroupSize)), 65535)", + "z": 1 }, "viewAlias": [ { "input": "past_k", "output": "present_key" }, @@ -6511,23 +5347,22 @@ "id": "present_append", "name": "GroupQueryAttention.MergeShareAppend", "shader": "gqa-present.wgsl.jinja", - "bindings": "presentMergeShareAppend", - "constants": { "shareRegion": "\"append\"" }, + "derive": { "shareRegion": "\"append\"" }, + "bindings": ["new_k", "new_v", "present_key_2", "present_value_2", "seqlens_k", "params_15"], "dispatch": { - "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)", - "workgroupSize": "copyWorkgroupSize" + "x": "min(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)), (copyWorkgroupSize)), 65535)", + "y": "ceilDiv(ceilDiv((dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.keyT, 1)), (copyWorkgroupSize)), 65535)", + "z": 1 } }, { "id": "attention", "name": "GroupQueryAttention.FlashPrefill", - "source": { - "shader": "attn-flash-prefill-cluster.wgsl.jinja", - "inputs": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" } - }, - "bindings": "flashPrefillSeqlens", + "shader": "attn-flash-prefill-cluster.wgsl.jinja", + "derive": { "layout": "\"bsh\"", "qLayout": "\"bsh\"", "kvLayout": "\"bhsd\"" }, + "bindings": ["query_8", "key_7", "value_7", "output_5", "seqlens_k", "params_28"], "dispatch": { - "x": "ceilDiv(dim(shapes.queryT, 1), constants.TILE_Q)", + "x": "ceilDiv(dim(shapes.queryT, 1), TILE_Q)", "y": "attrs.num_heads", "z": "dim(shapes.queryT, 0)" }