sync 6fdf6301e2bb

Browse files

Files changed (8) hide show

README.md +2 -2
build/webgpu/attn-flash-decode-splitk-merge.wgsl.jinja +2 -30
build/webgpu/bench.json +95 -18
build/webgpu/manifest.json +41 -49
build/webgpu/metadata.json +9 -9
build/webgpu/paged-attention.wgsl.jinja +5 -68
build/webgpu/paged-scatter-kv.wgsl.jinja +8 -38
build/webgpu/test.json +336 -1

README.md CHANGED Viewed

@@ -74,7 +74,7 @@ Every implementation variant requires `shader-f16`; the package has no variant-l
 - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
-- [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
 - [`attn-flash-decode-splitk-merge.wgsl.jinja`](build/webgpu/attn-flash-decode-splitk-merge.wgsl.jinja)
 - [`paged-attention.wgsl.jinja`](build/webgpu/paged-attention.wgsl.jinja)
 - [`paged-scatter-kv.wgsl.jinja`](build/webgpu/paged-scatter-kv.wgsl.jinja)
@@ -82,7 +82,7 @@ Every implementation variant requires `shader-f16`; the package has no variant-l
 ## Use with `@huggingface/kernels`
 ```sh
-npm install --save-exact @huggingface/kernels@0.0.1-preview.2
 ```
 Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.

 - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
+- [`bench.json`](build/webgpu/bench.json) — benchmark cases
 - [`attn-flash-decode-splitk-merge.wgsl.jinja`](build/webgpu/attn-flash-decode-splitk-merge.wgsl.jinja)
 - [`paged-attention.wgsl.jinja`](build/webgpu/paged-attention.wgsl.jinja)
 - [`paged-scatter-kv.wgsl.jinja`](build/webgpu/paged-scatter-kv.wgsl.jinja)
 ## Use with `@huggingface/kernels`
 ```sh
+npm install --save-exact @huggingface/kernels@0.0.1-preview.3
 ```
 Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.

build/webgpu/attn-flash-decode-splitk-merge.wgsl.jinja CHANGED Viewed

@@ -1,9 +1,3 @@
-{% if scalar == "f16" %}
-enable f16;
-{% endif %}
-{% if splitQueries is not defined %}{% set splitQueries = false %}{% endif %}
-{% if hasGate is not defined %}{% set hasGate = false %}{% endif %}
-{% set qSeq = qSeq | default(0) %}
 {{ env.wgsl.resourceDeclarations }}
 // Split-K flash decode, pass 2 of 2. Combines the per-split un-normalized online
@@ -19,9 +13,6 @@ enable f16;
 const HEAD_DIM_V4: u32 = {{ headDimV4 }}u;
 const Q_HEADS: u32 = {{ qNumHeads }}u;
 const NUM_SPLITS: u32 = {{ numSplits }}u;
-{% if splitQueries %}
-const Q_SEQ: u32 = {{ qSeq }}u;
-{% endif %}
 const Q_HIDDEN_V4: u32 = {{ qHiddenV4 }}u;
 // FLT_MAX, not -inf, as the online (m, d) accumulator init: merges must keep
 // `m - m` finite so an empty lane / all--inf row contributes the exact
@@ -42,6 +33,7 @@ fn shifted_value(value: f32, maxValue: f32) -> f32 {
   let equalFiniteMax = select(false, value == maxValue, is_finite_f32(maxValue));
   return select(value - maxValue, 0.0, equalFiniteMax);
 }
 fn exp_shift(value: f32, maxValue: f32) -> f32 {
   return exp(shifted_value(value, maxValue));
 }
@@ -50,21 +42,14 @@ fn main(
   @builtin(workgroup_id) wg: vec3<u32>,
   @builtin(local_invocation_id) lid: vec3<u32>
 ) {
-{% if splitQueries %}
-  let queryToken = wg.x;
-{% endif %}
   let h = wg.y;
   let b = wg.z;
   let d4 = lid.x;
-  if (h >= Q_HEADS || d4 >= HEAD_DIM_V4{% if splitQueries %} || queryToken >= Q_SEQ{% endif %}) {
     return;
   }
-{% if splitQueries %}
-  let mdBase = ((b * Q_SEQ + queryToken) * Q_HEADS + h) * NUM_SPLITS;
-{% else %}
   let mdBase = (b * Q_HEADS + h) * NUM_SPLITS;
-{% endif %}
   var globalMax = -FLT_MAX;
   for (var s: u32 = 0u; s < NUM_SPLITS; s = s + 1u) {
     globalMax = max(globalMax, partial_stats[mdBase + s].x);
@@ -75,28 +60,15 @@ fn main(
     let stats = partial_stats[mdBase + s];
     let w = exp_shift(stats.x, globalMax);
     globalDenom = globalDenom + stats.y * w;
-{% if splitQueries %}
-    let pBase = (((b * Q_SEQ + queryToken) * Q_HEADS + h) * NUM_SPLITS + s) * HEAD_DIM_V4;
-{% else %}
     let pBase = ((b * Q_HEADS + h) * NUM_SPLITS + s) * HEAD_DIM_V4;
-{% endif %}
     outV = outV + partial_out[pBase + d4] * w;
   }
-{% if splitQueries %}
-  let qBaseV4 = (b * Q_SEQ + queryToken) * Q_HIDDEN_V4 + h * HEAD_DIM_V4;
-{% else %}
   let qBaseV4 = b * Q_HIDDEN_V4 + h * HEAD_DIM_V4;
-{% endif %}
   // A split that covers no keys writes vec2(-FLT_MAX, 0), so when every split is empty —
   // a query whose window admits nothing — globalMax stays -FLT_MAX, exp_shift(x, x) is 1,
   // and globalDenom lands on exactly 0. Emit 0 rather than dividing so an all-empty
   // row has a defined zero result.
   var outValue = select(vec4<f32>(0.0), outV / globalDenom, globalDenom > 0.0);
-{% if hasGate %}
-  // The gated route multiplies the normalized attention output elementwise by its gate.
-  let gateV = vec4<f32>(gate[qBaseV4 + d4]);
-  outValue = outValue * (vec4<f32>(1.0) / (vec4<f32>(1.0) + exp(-gateV)));
-{% endif %}
   output[qBaseV4 + d4] = vec4<{{ scalar }}>(outValue);
 }

 {{ env.wgsl.resourceDeclarations }}
 // Split-K flash decode, pass 2 of 2. Combines the per-split un-normalized online
 const HEAD_DIM_V4: u32 = {{ headDimV4 }}u;
 const Q_HEADS: u32 = {{ qNumHeads }}u;
 const NUM_SPLITS: u32 = {{ numSplits }}u;
 const Q_HIDDEN_V4: u32 = {{ qHiddenV4 }}u;
 // FLT_MAX, not -inf, as the online (m, d) accumulator init: merges must keep
 // `m - m` finite so an empty lane / all--inf row contributes the exact
   let equalFiniteMax = select(false, value == maxValue, is_finite_f32(maxValue));
   return select(value - maxValue, 0.0, equalFiniteMax);
 }
 fn exp_shift(value: f32, maxValue: f32) -> f32 {
   return exp(shifted_value(value, maxValue));
 }
   @builtin(workgroup_id) wg: vec3<u32>,
   @builtin(local_invocation_id) lid: vec3<u32>
 ) {
   let h = wg.y;
   let b = wg.z;
   let d4 = lid.x;
+  if (h >= Q_HEADS || d4 >= HEAD_DIM_V4) {
     return;
   }
   let mdBase = (b * Q_HEADS + h) * NUM_SPLITS;
   var globalMax = -FLT_MAX;
   for (var s: u32 = 0u; s < NUM_SPLITS; s = s + 1u) {
     globalMax = max(globalMax, partial_stats[mdBase + s].x);
     let stats = partial_stats[mdBase + s];
     let w = exp_shift(stats.x, globalMax);
     globalDenom = globalDenom + stats.y * w;
     let pBase = ((b * Q_HEADS + h) * NUM_SPLITS + s) * HEAD_DIM_V4;
     outV = outV + partial_out[pBase + d4] * w;
   }
   let qBaseV4 = b * Q_HIDDEN_V4 + h * HEAD_DIM_V4;
   // A split that covers no keys writes vec2(-FLT_MAX, 0), so when every split is empty —
   // a query whose window admits nothing — globalMax stays -FLT_MAX, exp_shift(x, x) is 1,
   // and globalDenom lands on exactly 0. Emit 0 rather than dividing so an all-empty
   // row has a defined zero result.
   var outValue = select(vec4<f32>(0.0), outV / globalDenom, globalDenom > 0.0);
   output[qBaseV4 + d4] = vec4<{{ scalar }}>(outValue);
 }

build/webgpu/bench.json CHANGED Viewed

@@ -38,7 +38,7 @@
       "name": "paged-decode-llama-b16-past2048-h32kv32-d128",
       "preset": "model",
       "provenance": {
-        "notes": "Llama class defaults (32 query heads, no GQA, head_dim 128) decoding 16 concurrent sequences with 2048 tokens of paged KV each."
       },
       "vars": { "dtype": "float16", "batch": 16, "past": 2048, "heads": 32, "kvHeads": 32, "headDim": 128 },
       "attrs": { "num_heads": 32, "kv_num_heads": 32 },
@@ -57,12 +57,9 @@
         "cumulativeSequenceLengthT": {
           "shape": [17],
           "dtype": "int32",
-          "dist": "linearMod",
-          "seed": 7605,
-          "step": 1,
-          "mod": 17
         },
-        "pastSeqlensT": { "shape": [16], "dtype": "int32", "dist": "constant", "value": 2048 },
         "blockTableT": {
           "shape": [16, 128],
           "dtype": "int32",
@@ -81,7 +78,7 @@
       "name": "paged-decode-mistral-b16-past4096-h32kv8-d128",
       "preset": "model",
       "provenance": {
-        "notes": "Mistral class defaults (32 query heads, 8 key/value heads, head_dim 128) at 4096 tokens of history -- the GQA case, where the cache is a quarter the size for the same context."
       },
       "vars": { "dtype": "float16", "batch": 16, "past": 4096, "heads": 32, "kvHeads": 8, "headDim": 128 },
       "attrs": { "num_heads": 32, "kv_num_heads": 8 },
@@ -94,12 +91,9 @@
         "cumulativeSequenceLengthT": {
           "shape": [17],
           "dtype": "int32",
-          "dist": "linearMod",
-          "seed": 7705,
-          "step": 1,
-          "mod": 17
         },
-        "pastSeqlensT": { "shape": [16], "dtype": "int32", "dist": "constant", "value": 4096 },
         "blockTableT": {
           "shape": [16, 256],
           "dtype": "int32",
@@ -118,7 +112,7 @@
       "name": "paged-decode-gemma3-b8-past2048-h8kv4-d256",
       "preset": "model",
       "provenance": {
-        "notes": "Gemma3 class defaults (8 query heads, 4 key/value heads, head_dim 256) -- few heads but the widest head dimension in this set."
       },
       "vars": { "dtype": "float16", "batch": 8, "past": 2048, "heads": 8, "kvHeads": 4, "headDim": 256 },
       "attrs": { "num_heads": 8, "kv_num_heads": 4 },
@@ -131,12 +125,9 @@
         "cumulativeSequenceLengthT": {
           "shape": [9],
           "dtype": "int32",
-          "dist": "linearMod",
-          "seed": 7805,
-          "step": 1,
-          "mod": 9
         },
-        "pastSeqlensT": { "shape": [8], "dtype": "int32", "dist": "constant", "value": 2048 },
         "blockTableT": {
           "shape": [8, 128],
           "dtype": "int32",
@@ -150,6 +141,92 @@
       "bench": {
         "metrics": [{ "type": "bandwidth", "value": "2 * args.batch * args.kvHeads * args.past * args.headDim * 2" }]
       }
     }
   ]
 }

       "name": "paged-decode-llama-b16-past2048-h32kv32-d128",
       "preset": "model",
       "provenance": {
+        "notes": "Llama class defaults (32 query heads, no GQA, head_dim 128) decoding 16 concurrent sequences with 2048 tokens of paged KV each. Context counts include this step: past_seqlens is context - 1 and the exclusive prefix [0, ..., batch] assigns one new token per sequence. All logical blocks map to distinct physical blocks, including the final append slot."
       },
       "vars": { "dtype": "float16", "batch": 16, "past": 2048, "heads": 32, "kvHeads": 32, "headDim": 128 },
       "attrs": { "num_heads": 32, "kv_num_heads": 32 },
         "cumulativeSequenceLengthT": {
           "shape": [17],
           "dtype": "int32",
+          "data": { "kind": "linspace", "start": 0, "end": 16 }
         },
+        "pastSeqlensT": { "shape": [16], "dtype": "int32", "dist": "constant", "value": 2047 },
         "blockTableT": {
           "shape": [16, 128],
           "dtype": "int32",
       "name": "paged-decode-mistral-b16-past4096-h32kv8-d128",
       "preset": "model",
       "provenance": {
+        "notes": "Mistral class defaults (32 query heads, 8 key/value heads, head_dim 128) at 4096 tokens per sequence including the current token -- the GQA case, where the cache is a quarter the size for the same context. Context counts include this step: past_seqlens is context - 1 and the exclusive prefix [0, ..., batch] assigns one new token per sequence. All logical blocks map to distinct physical blocks, including the final append slot."
       },
       "vars": { "dtype": "float16", "batch": 16, "past": 4096, "heads": 32, "kvHeads": 8, "headDim": 128 },
       "attrs": { "num_heads": 32, "kv_num_heads": 8 },
         "cumulativeSequenceLengthT": {
           "shape": [17],
           "dtype": "int32",
+          "data": { "kind": "linspace", "start": 0, "end": 16 }
         },
+        "pastSeqlensT": { "shape": [16], "dtype": "int32", "dist": "constant", "value": 4095 },
         "blockTableT": {
           "shape": [16, 256],
           "dtype": "int32",
       "name": "paged-decode-gemma3-b8-past2048-h8kv4-d256",
       "preset": "model",
       "provenance": {
+        "notes": "Gemma3 class defaults (8 query heads, 4 key/value heads, head_dim 256) -- few heads but the widest head dimension in this set. Context counts include this step: past_seqlens is context - 1 and the exclusive prefix [0, ..., batch] assigns one new token per sequence. All logical blocks map to distinct physical blocks, including the final append slot."
       },
       "vars": { "dtype": "float16", "batch": 8, "past": 2048, "heads": 8, "kvHeads": 4, "headDim": 256 },
       "attrs": { "num_heads": 8, "kv_num_heads": 4 },
         "cumulativeSequenceLengthT": {
           "shape": [9],
           "dtype": "int32",
+          "data": { "kind": "linspace", "start": 0, "end": 8 }
         },
+        "pastSeqlensT": { "shape": [8], "dtype": "int32", "dist": "constant", "value": 2047 },
         "blockTableT": {
           "shape": [8, 128],
           "dtype": "int32",
       "bench": {
         "metrics": [{ "type": "bandwidth", "value": "2 * args.batch * args.kvHeads * args.past * args.headDim * 2" }]
       }
+    },
+    {
+      "name": "scatter-fold-wg1-independent-sequences",
+      "provenance": {
+        "notes": "A large scatter workload crosses the WebGPU dispatch-grid limit. One-token sequences bound attention cost; nonzero dyadic cycles expose missing tail writes and misplaced cache/value elements."
+      },
+      "attrs": { "num_heads": 8, "kv_num_heads": 8 },
+      "tunables": { "WORKGROUP_SIZE": 64, "SCATTER_WORKGROUP_SIZE": 1 },
+      "inputs": {
+        "queryT": { "dtype": "float16", "shape": [64, 1024], "data": { "kind": "constant", "value": 0.0 } },
+        "keyT": {
+          "dtype": "float16",
+          "shape": [64, 1024],
+          "data": { "kind": "cycle", "values": [1.0, 2.0, 4.0, -1.0, -2.0] }
+        },
+        "valueT": {
+          "dtype": "float16",
+          "shape": [64, 1024],
+          "data": { "kind": "cycle", "values": [2.0, -4.0, 8.0, 1.0, -2.0, 4.0, -1.0] }
+        },
+        "keyCacheT": { "dtype": "float16", "shape": [64, 1, 8, 128], "data": { "kind": "constant", "value": 0.0 } },
+        "valueCacheT": { "dtype": "float16", "shape": [64, 1, 8, 128], "data": { "kind": "constant", "value": 0.0 } },
+        "cumulativeSequenceLengthT": {
+          "dtype": "int32",
+          "shape": [65],
+          "data": { "kind": "linspace", "start": 0, "end": 64 }
+        },
+        "pastSeqlensT": { "dtype": "int32", "shape": [64], "data": { "kind": "constant", "value": 0 } },
+        "blockTableT": { "dtype": "int32", "shape": [64, 1], "data": { "kind": "linspace", "start": 0, "end": 63 } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float16", "shape": [64, 1024] },
+        "keyCacheT": { "dtype": "float16", "shape": [64, 1, 8, 128] },
+        "valueCacheT": { "dtype": "float16", "shape": [64, 1, 8, 128] }
+      },
+      "preset": "edge",
+      "vars": { "tokens": 64, "heads": 8, "headDim": 128 },
+      "tunableSpace": { "WORKGROUP_SIZE": [64], "SCATTER_WORKGROUP_SIZE": [1] },
+      "bench": {
+        "metrics": [
+          { "type": "bandwidth", "value": "6 * args.tokens * args.heads * args.headDim * 2", "provenance": "logical" }
+        ]
+      }
+    },
+    {
+      "name": "scatter-fold-wg4-independent-sequences",
+      "provenance": {
+        "notes": "A large scatter workload crosses the WebGPU dispatch-grid limit. One-token sequences bound attention cost; nonzero dyadic cycles expose missing tail writes and misplaced cache/value elements."
+      },
+      "attrs": { "num_heads": 8, "kv_num_heads": 8 },
+      "tunables": { "WORKGROUP_SIZE": 64, "SCATTER_WORKGROUP_SIZE": 4 },
+      "inputs": {
+        "queryT": { "dtype": "float16", "shape": [256, 1024], "data": { "kind": "constant", "value": 0.0 } },
+        "keyT": {
+          "dtype": "float16",
+          "shape": [256, 1024],
+          "data": { "kind": "cycle", "values": [1.0, 2.0, 4.0, -1.0, -2.0] }
+        },
+        "valueT": {
+          "dtype": "float16",
+          "shape": [256, 1024],
+          "data": { "kind": "cycle", "values": [2.0, -4.0, 8.0, 1.0, -2.0, 4.0, -1.0] }
+        },
+        "keyCacheT": { "dtype": "float16", "shape": [256, 1, 8, 128], "data": { "kind": "constant", "value": 0.0 } },
+        "valueCacheT": { "dtype": "float16", "shape": [256, 1, 8, 128], "data": { "kind": "constant", "value": 0.0 } },
+        "cumulativeSequenceLengthT": {
+          "dtype": "int32",
+          "shape": [257],
+          "data": { "kind": "linspace", "start": 0, "end": 256 }
+        },
+        "pastSeqlensT": { "dtype": "int32", "shape": [256], "data": { "kind": "constant", "value": 0 } },
+        "blockTableT": { "dtype": "int32", "shape": [256, 1], "data": { "kind": "linspace", "start": 0, "end": 255 } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float16", "shape": [256, 1024] },
+        "keyCacheT": { "dtype": "float16", "shape": [256, 1, 8, 128] },
+        "valueCacheT": { "dtype": "float16", "shape": [256, 1, 8, 128] }
+      },
+      "preset": "edge",
+      "vars": { "tokens": 256, "heads": 8, "headDim": 128 },
+      "tunableSpace": { "WORKGROUP_SIZE": [64], "SCATTER_WORKGROUP_SIZE": [4] },
+      "bench": {
+        "metrics": [
+          { "type": "bandwidth", "value": "6 * args.tokens * args.heads * args.headDim * 2", "provenance": "logical" }
+        ]
+      }
     }
   ]
 }

build/webgpu/manifest.json CHANGED Viewed

@@ -86,51 +86,43 @@
     "headDim": "headSize",
     "qNumHeads": "attrs.num_heads",
     "qHiddenV4": "qHidden / 4",
-    "hasBias": false,
     "outVec": "\"vec4<f16>\" if dtypes.T == \"f16\" else \"vec4<f32>\""
   },
   "when": ["pagedContractOk", "dispatchFits"],
   "bindings": {
-    "key": { "arg": "keyT", "buffer": "read-only-storage", "elementType": "$aScalar" },
-    "value": { "arg": "valueT", "buffer": "read-only-storage", "elementType": "$aScalar" },
-    "key_cache": { "arg": "keyCacheT", "buffer": "storage", "elementType": "$aScalar" },
-    "value_cache": { "arg": "valueCacheT", "buffer": "storage", "elementType": "$aScalar" },
-    "cumulative_sequence_length": {
-      "arg": "cumulativeSequenceLengthT",
-      "buffer": "read-only-storage",
-      "elementType": "i32"
-    },
-    "past_seqlens": { "arg": "pastSeqlensT", "buffer": "read-only-storage", "elementType": "i32" },
-    "block_table": { "arg": "blockTableT", "buffer": "read-only-storage", "elementType": "i32" },
     "params": {
-      "buffer": "uniform",
       "struct": [
         { "name": "batchSize", "type": "u32", "value": "batchSize" },
         { "name": "scatterCount", "type": "u32", "value": "tokenCount * attrs.kv_num_heads * headSize" }
       ]
     },
-    "slot_mapping": { "arg": "slotMappingT", "buffer": "read-only-storage", "elementType": "i32" },
-    "params_2": {
       "name": "params",
-      "buffer": "uniform",
       "struct": [{ "name": "scatterCount", "type": "u32", "value": "tokenCount * attrs.kv_num_heads * headSize" }]
     },
-    "query": { "arg": "queryT", "buffer": "read-only-storage", "elementType": "$aScalar" },
-    "key_cache_2": {
       "arg": "keyCacheT",
       "name": "key_cache",
       "buffer": "read-only-storage",
       "elementType": "$cacheVec"
     },
-    "value_cache_2": {
       "arg": "valueCacheT",
       "name": "value_cache",
       "buffer": "read-only-storage",
       "elementType": "$cacheVec"
     },
-    "params_3": {
       "name": "params",
-      "buffer": "uniform",
       "struct": [
         { "name": "batchSize", "type": "u32", "value": "batchSize" },
         { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" }
@@ -167,14 +159,14 @@
           "derive": { "splitK": true },
           "bindings": [
             "query",
-            "key_cache_2",
-            "value_cache_2",
             "cumulative_sequence_length",
             "past_seqlens",
             "block_table",
             { "scratch": "partialOut", "name": "partial_out", "elementType": "vec4<f32>" },
             { "scratch": "partialStats", "name": "partial_stats", "elementType": "vec2<f32>" },
-            "params_3"
           ],
           "dispatch": { "x": "tokenCount", "y": "attrs.kv_num_heads", "z": "numSplits" }
         },
@@ -224,13 +216,13 @@
           "shader": "paged-attention.wgsl.jinja",
           "bindings": [
             "query",
-            "key_cache_2",
-            "value_cache_2",
             "cumulative_sequence_length",
             "past_seqlens",
             "block_table",
             { "arg": "outputT", "name": "output", "elementType": "$aScalar" },
-            "params_3"
           ],
           "dispatch": { "x": "tokenCount", "y": "attrs.kv_num_heads" }
         }
@@ -251,7 +243,7 @@
           "id": "scatter",
           "name": "PagedAttention.ScatterKV",
           "shader": "paged-scatter-kv.wgsl.jinja",
-          "bindings": ["key", "value", "key_cache", "value_cache", "slot_mapping", "params_2"],
           "dispatch": {
             "x": "min(ceilDiv((tokenCount * attrs.kv_num_heads * headSize), (tunables.SCATTER_WORKGROUP_SIZE)), 65535)",
             "y": "ceilDiv(ceilDiv((tokenCount * attrs.kv_num_heads * headSize), (tunables.SCATTER_WORKGROUP_SIZE)), 65535)",
@@ -265,14 +257,14 @@
           "derive": { "splitK": true },
           "bindings": [
             "query",
-            "key_cache_2",
-            "value_cache_2",
             "cumulative_sequence_length",
             "past_seqlens",
             "block_table",
             { "scratch": "partialOut", "name": "partial_out", "elementType": "vec4<f32>" },
             { "scratch": "partialStats", "name": "partial_stats", "elementType": "vec2<f32>" },
-            "params_3"
           ],
           "dispatch": { "x": "tokenCount", "y": "attrs.kv_num_heads", "z": "numSplits" }
         },
@@ -309,7 +301,7 @@
           "id": "scatter",
           "name": "PagedAttention.ScatterKV",
           "shader": "paged-scatter-kv.wgsl.jinja",
-          "bindings": ["key", "value", "key_cache", "value_cache", "slot_mapping", "params_2"],
           "dispatch": {
             "x": "min(ceilDiv((tokenCount * attrs.kv_num_heads * headSize), (tunables.SCATTER_WORKGROUP_SIZE)), 65535)",
             "y": "ceilDiv(ceilDiv((tokenCount * attrs.kv_num_heads * headSize), (tunables.SCATTER_WORKGROUP_SIZE)), 65535)",
@@ -322,13 +314,13 @@
           "shader": "paged-attention.wgsl.jinja",
           "bindings": [
             "query",
-            "key_cache_2",
-            "value_cache_2",
             "cumulative_sequence_length",
             "past_seqlens",
             "block_table",
             { "arg": "outputT", "name": "output", "elementType": "$aScalar" },
-            "params_3"
           ],
           "dispatch": { "x": "tokenCount", "y": "attrs.kv_num_heads" }
         }
@@ -363,14 +355,14 @@
           "derive": { "splitK": true },
           "bindings": [
             "query",
-            "key_cache_2",
-            "value_cache_2",
             "cumulative_sequence_length",
             "past_seqlens",
             "block_table",
             { "scratch": "partialOut", "name": "partial_out", "elementType": "vec4<f32>" },
             { "scratch": "partialStats", "name": "partial_stats", "elementType": "vec2<f32>" },
-            "params_3"
           ],
           "dispatch": { "x": "tokenCount", "y": "attrs.kv_num_heads", "z": "numSplits" }
         },
@@ -420,13 +412,13 @@
           "shader": "paged-attention.wgsl.jinja",
           "bindings": [
             "query",
-            "key_cache_2",
-            "value_cache_2",
             "cumulative_sequence_length",
             "past_seqlens",
             "block_table",
             { "arg": "outputT", "name": "output", "elementType": "$aScalar" },
-            "params_3"
           ],
           "dispatch": { "x": "tokenCount", "y": "attrs.kv_num_heads" }
         }
@@ -447,7 +439,7 @@
           "id": "scatter",
           "name": "PagedAttention.ScatterKV",
           "shader": "paged-scatter-kv.wgsl.jinja",
-          "bindings": ["query", "key_cache", "value_cache", "slot_mapping", "params_2"],
           "dispatch": {
             "x": "min(ceilDiv((tokenCount * attrs.kv_num_heads * headSize), (tunables.SCATTER_WORKGROUP_SIZE)), 65535)",
             "y": "ceilDiv(ceilDiv((tokenCount * attrs.kv_num_heads * headSize), (tunables.SCATTER_WORKGROUP_SIZE)), 65535)",
@@ -461,14 +453,14 @@
           "derive": { "splitK": true },
           "bindings": [
             "query",
-            "key_cache_2",
-            "value_cache_2",
             "cumulative_sequence_length",
             "past_seqlens",
             "block_table",
             { "scratch": "partialOut", "name": "partial_out", "elementType": "vec4<f32>" },
             { "scratch": "partialStats", "name": "partial_stats", "elementType": "vec2<f32>" },
-            "params_3"
           ],
           "dispatch": { "x": "tokenCount", "y": "attrs.kv_num_heads", "z": "numSplits" }
         },
@@ -505,7 +497,7 @@
           "id": "scatter",
           "name": "PagedAttention.ScatterKV",
           "shader": "paged-scatter-kv.wgsl.jinja",
-          "bindings": ["query", "key_cache", "value_cache", "slot_mapping", "params_2"],
           "dispatch": {
             "x": "min(ceilDiv((tokenCount * attrs.kv_num_heads * headSize), (tunables.SCATTER_WORKGROUP_SIZE)), 65535)",
             "y": "ceilDiv(ceilDiv((tokenCount * attrs.kv_num_heads * headSize), (tunables.SCATTER_WORKGROUP_SIZE)), 65535)",
@@ -518,13 +510,13 @@
           "shader": "paged-attention.wgsl.jinja",
           "bindings": [
             "query",
-            "key_cache_2",
-            "value_cache_2",
             "cumulative_sequence_length",
             "past_seqlens",
             "block_table",
             { "arg": "outputT", "name": "output", "elementType": "$aScalar" },
-            "params_3"
           ],
           "dispatch": { "x": "tokenCount", "y": "attrs.kv_num_heads" }
         }

     "headDim": "headSize",
     "qNumHeads": "attrs.num_heads",
     "qHiddenV4": "qHidden / 4",
     "outVec": "\"vec4<f16>\" if dtypes.T == \"f16\" else \"vec4<f32>\""
   },
   "when": ["pagedContractOk", "dispatchFits"],
   "bindings": {
+    "key": { "arg": "keyT", "elementType": "$aScalar" },
+    "value": { "arg": "valueT", "elementType": "$aScalar" },
+    "key_cache": { "arg": "keyCacheT", "elementType": "$aScalar" },
+    "value_cache": { "arg": "valueCacheT", "elementType": "$aScalar" },
+    "cumulative_sequence_length": { "arg": "cumulativeSequenceLengthT", "elementType": "i32" },
+    "past_seqlens": { "arg": "pastSeqlensT", "elementType": "i32" },
+    "block_table": { "arg": "blockTableT", "elementType": "i32" },
     "params": {
       "struct": [
         { "name": "batchSize", "type": "u32", "value": "batchSize" },
         { "name": "scatterCount", "type": "u32", "value": "tokenCount * attrs.kv_num_heads * headSize" }
       ]
     },
+    "slot_mapping": { "arg": "slotMappingT", "elementType": "i32" },
+    "params_scatter": {
       "name": "params",
       "struct": [{ "name": "scatterCount", "type": "u32", "value": "tokenCount * attrs.kv_num_heads * headSize" }]
     },
+    "query": { "arg": "queryT", "elementType": "$aScalar" },
+    "key_cache_main": {
       "arg": "keyCacheT",
       "name": "key_cache",
       "buffer": "read-only-storage",
       "elementType": "$cacheVec"
     },
+    "value_cache_main": {
       "arg": "valueCacheT",
       "name": "value_cache",
       "buffer": "read-only-storage",
       "elementType": "$cacheVec"
     },
+    "params_main": {
       "name": "params",
       "struct": [
         { "name": "batchSize", "type": "u32", "value": "batchSize" },
         { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" }
           "derive": { "splitK": true },
           "bindings": [
             "query",
+            "key_cache_main",
+            "value_cache_main",
             "cumulative_sequence_length",
             "past_seqlens",
             "block_table",
             { "scratch": "partialOut", "name": "partial_out", "elementType": "vec4<f32>" },
             { "scratch": "partialStats", "name": "partial_stats", "elementType": "vec2<f32>" },
+            "params_main"
           ],
           "dispatch": { "x": "tokenCount", "y": "attrs.kv_num_heads", "z": "numSplits" }
         },
           "shader": "paged-attention.wgsl.jinja",
           "bindings": [
             "query",
+            "key_cache_main",
+            "value_cache_main",
             "cumulative_sequence_length",
             "past_seqlens",
             "block_table",
             { "arg": "outputT", "name": "output", "elementType": "$aScalar" },
+            "params_main"
           ],
           "dispatch": { "x": "tokenCount", "y": "attrs.kv_num_heads" }
         }
           "id": "scatter",
           "name": "PagedAttention.ScatterKV",
           "shader": "paged-scatter-kv.wgsl.jinja",
+          "bindings": ["key", "value", "key_cache", "value_cache", "slot_mapping", "params_scatter"],
           "dispatch": {
             "x": "min(ceilDiv((tokenCount * attrs.kv_num_heads * headSize), (tunables.SCATTER_WORKGROUP_SIZE)), 65535)",
             "y": "ceilDiv(ceilDiv((tokenCount * attrs.kv_num_heads * headSize), (tunables.SCATTER_WORKGROUP_SIZE)), 65535)",
           "derive": { "splitK": true },
           "bindings": [
             "query",
+            "key_cache_main",
+            "value_cache_main",
             "cumulative_sequence_length",
             "past_seqlens",
             "block_table",
             { "scratch": "partialOut", "name": "partial_out", "elementType": "vec4<f32>" },
             { "scratch": "partialStats", "name": "partial_stats", "elementType": "vec2<f32>" },
+            "params_main"
           ],
           "dispatch": { "x": "tokenCount", "y": "attrs.kv_num_heads", "z": "numSplits" }
         },
           "id": "scatter",
           "name": "PagedAttention.ScatterKV",
           "shader": "paged-scatter-kv.wgsl.jinja",
+          "bindings": ["key", "value", "key_cache", "value_cache", "slot_mapping", "params_scatter"],
           "dispatch": {
             "x": "min(ceilDiv((tokenCount * attrs.kv_num_heads * headSize), (tunables.SCATTER_WORKGROUP_SIZE)), 65535)",
             "y": "ceilDiv(ceilDiv((tokenCount * attrs.kv_num_heads * headSize), (tunables.SCATTER_WORKGROUP_SIZE)), 65535)",
           "shader": "paged-attention.wgsl.jinja",
           "bindings": [
             "query",
+            "key_cache_main",
+            "value_cache_main",
             "cumulative_sequence_length",
             "past_seqlens",
             "block_table",
             { "arg": "outputT", "name": "output", "elementType": "$aScalar" },
+            "params_main"
           ],
           "dispatch": { "x": "tokenCount", "y": "attrs.kv_num_heads" }
         }
           "derive": { "splitK": true },
           "bindings": [
             "query",
+            "key_cache_main",
+            "value_cache_main",
             "cumulative_sequence_length",
             "past_seqlens",
             "block_table",
             { "scratch": "partialOut", "name": "partial_out", "elementType": "vec4<f32>" },
             { "scratch": "partialStats", "name": "partial_stats", "elementType": "vec2<f32>" },
+            "params_main"
           ],
           "dispatch": { "x": "tokenCount", "y": "attrs.kv_num_heads", "z": "numSplits" }
         },
           "shader": "paged-attention.wgsl.jinja",
           "bindings": [
             "query",
+            "key_cache_main",
+            "value_cache_main",
             "cumulative_sequence_length",
             "past_seqlens",
             "block_table",
             { "arg": "outputT", "name": "output", "elementType": "$aScalar" },
+            "params_main"
           ],
           "dispatch": { "x": "tokenCount", "y": "attrs.kv_num_heads" }
         }
           "id": "scatter",
           "name": "PagedAttention.ScatterKV",
           "shader": "paged-scatter-kv.wgsl.jinja",
+          "bindings": ["query", "key_cache", "value_cache", "slot_mapping", "params_scatter"],
           "dispatch": {
             "x": "min(ceilDiv((tokenCount * attrs.kv_num_heads * headSize), (tunables.SCATTER_WORKGROUP_SIZE)), 65535)",
             "y": "ceilDiv(ceilDiv((tokenCount * attrs.kv_num_heads * headSize), (tunables.SCATTER_WORKGROUP_SIZE)), 65535)",
           "derive": { "splitK": true },
           "bindings": [
             "query",
+            "key_cache_main",
+            "value_cache_main",
             "cumulative_sequence_length",
             "past_seqlens",
             "block_table",
             { "scratch": "partialOut", "name": "partial_out", "elementType": "vec4<f32>" },
             { "scratch": "partialStats", "name": "partial_stats", "elementType": "vec2<f32>" },
+            "params_main"
           ],
           "dispatch": { "x": "tokenCount", "y": "attrs.kv_num_heads", "z": "numSplits" }
         },
           "id": "scatter",
           "name": "PagedAttention.ScatterKV",
           "shader": "paged-scatter-kv.wgsl.jinja",
+          "bindings": ["query", "key_cache", "value_cache", "slot_mapping", "params_scatter"],
           "dispatch": {
             "x": "min(ceilDiv((tokenCount * attrs.kv_num_heads * headSize), (tunables.SCATTER_WORKGROUP_SIZE)), 65535)",
             "y": "ceilDiv(ceilDiv((tokenCount * attrs.kv_num_heads * headSize), (tunables.SCATTER_WORKGROUP_SIZE)), 65535)",
           "shader": "paged-attention.wgsl.jinja",
           "bindings": [
             "query",
+            "key_cache_main",
+            "value_cache_main",
             "cumulative_sequence_length",
             "past_seqlens",
             "block_table",
             { "arg": "outputT", "name": "output", "elementType": "$aScalar" },
+            "params_main"
           ],
           "dispatch": { "x": "tokenCount", "y": "attrs.kv_num_heads" }
         }

build/webgpu/metadata.json CHANGED Viewed

@@ -1,23 +1,23 @@
 {
   "name": "com.microsoft.PagedAttention",
-  "id": "_com_microsoft_pagedattention_webgpu_a50bcf6",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "attn-flash-decode-splitk-merge.wgsl.jinja": "uOgPgm5HioxH0ha2Z7Q0b/8qESsCmwKAS/198ogYX3I=",
-      "bench.json": "n1VW3Kl/whbn6CsIdfs0MHNrvUTskejhRkd5jqAAGmk=",
-      "manifest.json": "T7kLnpUD+K08/2E51GlZEK8v9YsGXliBIVnku1LRgRM=",
-      "paged-attention.wgsl.jinja": "NS4JduPGQ8x2YxgIap9fKuggGtyqse0u503A7rvzSM4=",
-      "paged-scatter-kv.wgsl.jinja": "EEb8eRJICdQOiQPsThu7rGm97W8iYFFH5zL3vArnQ8Q=",
-      "test.json": "UtJ6CP4M4fFlZm0lZdIrcX+8AxB/5gIC88fXZb6Iw5Q="
     }
   },
-  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
   "webgpu": {
-    "manifestSpec": "2.0",
     "variants": {
       "separate_derived_splitk": ["attn-flash-decode-splitk-merge.wgsl.jinja", "paged-attention.wgsl.jinja", "paged-scatter-kv.wgsl.jinja"],
       "separate_derived": ["paged-attention.wgsl.jinja", "paged-scatter-kv.wgsl.jinja"],

 {
   "name": "com.microsoft.PagedAttention",
+  "id": "_com_microsoft_pagedattention_webgpu_bcb1627",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "attn-flash-decode-splitk-merge.wgsl.jinja": "P0yN9VRWOBKM6lMTWKDX7eKuxcmVjj3YX+eRlYgyG0w=",
+      "bench.json": "wiVFqyO9vEyYhD/i3BuWEwEMtorPGtDpjrJQSTX+VEU=",
+      "manifest.json": "RVVRT3P/Ix1RKhWszk/4atE5P78+OEHtHAkSyDYe1v0=",
+      "paged-attention.wgsl.jinja": "Tkqnp3/PxodwRygVuawTRrPNrS2WDP+n/IYU6693Hog=",
+      "paged-scatter-kv.wgsl.jinja": "Xq8l6iRmNLXlx510ylhkcIj4AQSbb5jM5M0/A1IOp7Q=",
+      "test.json": "zGwuobgxFob47dayYp0kP1EbTZE9CSCKd0pFZjptm40="
     }
   },
+  "provenance": { "kernel": { "sha": "6fdf6301e2bbcc2f03bf1eaf493b7ad55ef33afc", "dirty": false } },
   "webgpu": {
+    "manifestSpec": "2.1",
     "variants": {
       "separate_derived_splitk": ["attn-flash-decode-splitk-merge.wgsl.jinja", "paged-attention.wgsl.jinja", "paged-scatter-kv.wgsl.jinja"],
       "separate_derived": ["paged-attention.wgsl.jinja", "paged-scatter-kv.wgsl.jinja"],

build/webgpu/paged-attention.wgsl.jinja CHANGED Viewed

@@ -1,7 +1,6 @@
 {% set splitK = splitK is defined and splitK %}
 {% set keyLo = "splitStart" if splitK else "0u" %}
 {% set keyHi = "splitEnd" if splitK else "key_bound" %}
-enable f16;
 {{ env.wgsl.resourceDeclarations }}
 // com.microsoft.PagedAttention, attention pass.
@@ -56,6 +55,7 @@ fn shifted_value(value: f32, maxValue: f32) -> f32 {
   let equalFiniteMax = select(false, value == maxValue, is_finite_f32(maxValue));
   return select(value - maxValue, 0.0, equalFiniteMax);
 }
 fn exp_shift(value: f32, maxValue: f32) -> f32 {
   return exp(shifted_value(value, maxValue));
 }
@@ -78,59 +78,6 @@ var<workgroup> key_rows: array<u32, WG>;
 // Both the subgroup and portable barrier-tree engines return the same merged
 // pair to every invocation. Repeated merges require a workgroup barrier between
 // calls before their shared partial storage is reused.
-{% set combineSubgroups = combineSubgroups is defined and combineSubgroups %}
-{% if combineSubgroups %}
-// Cross-subgroup merge that assumes nothing about which invocations share a
-// subgroup or how many subgroups there are: each subgroup's elected lane
-// publishes the subgroup pair in the slot at its OWN invocation index and sets
-// that index's bit in a workgroup bitmask; thread 0 then folds exactly the
-// published slots, in ascending index order (the online (m, d) merge is not
-// float-associative, so the order is fixed), and clears the mask for the next
-// call as it reads it. Workgroup memory starts zeroed, so the mask needs no
-// setup. Same three collectives as a single-subgroup reduce, two barriers.
-var<workgroup> partialM: array<f32, WG>;
-var<workgroup> partialD: array<f32, WG>;
-var<workgroup> leaderMask: array<atomic<u32>, (WG + 31u) / 32u>;
-var<workgroup> combinedMD: vec2<f32>;
-// When the whole workgroup is one subgroup the subgroup reduce already covers
-// it (no barriers, no shared state). `subgroup_size` is the size of the current
-// subgroup and uniform, so the test is exact and may guard the barriers below.
-fn combine_partials(m: f32, d: f32, lidx: u32, sgSize: u32) -> vec2<f32> {
-  let sgM = subgroupMax(m);
-  // A lane with no elements contributes d == 0 (exact identity). A +inf
-  // element made exp(inf - inf) = NaN stick in that lane's d; a NaN element
-  // landed in d via exp(NaN); both survive the merge and are detected by the
-  // code after the reduction.
-  let sgD = subgroupAdd(d * exp_shift(m, sgM));
-  if (sgSize == WG) {
-    return vec2<f32>(sgM, sgD);
-  }
-  if (subgroupElect()) {
-    partialM[lidx] = sgM;
-    partialD[lidx] = sgD;
-    atomicOr(&leaderMask[lidx / 32u], 1u << (lidx % 32u));
-  }
-  workgroupBarrier();
-  if (lidx == 0u) {
-    var accM = -FLT_MAX;
-    var accD = 0.0;
-    for (var w = 0u; w < (WG + 31u) / 32u; w = w + 1u) {
-      var bits = atomicExchange(&leaderMask[w], 0u);
-      while (bits != 0u) {
-        let slot = w * 32u + firstTrailingBit(bits);
-        bits = bits & (bits - 1u);
-        let mNew = max(accM, partialM[slot]);
-        accD = accD * exp_shift(accM, mNew) + partialD[slot] * exp_shift(partialM[slot], mNew);
-        accM = mNew;
-      }
-    }
-    combinedMD = vec2<f32>(accM, accD);
-  }
-  workgroupBarrier();
-  return combinedMD;
-}
-{% else %}
 {% set mdStreamed = mdStreams is defined %}
 {% set mdStreams = mdStreams if mdStreams is defined else 1 %}
 {% set mdExtent = "WG" if mdStreams == 1 else "WG * " ~ mdStreams ~ "u" %}
@@ -195,16 +142,12 @@ fn combine_partials(m: f32, d: f32, lidx: u32) -> vec2<f32> {
   return merged;
 }
 {% endif %}
-{% endif %}
-{% if ATTN_SCALE_DIM is not defined %}{% set ATTN_SCALE_DIM = "HEAD_DIM" %}{% endif %}
-fn scale_value() -> f32 {
   if (params.scale != 0.0) { return params.scale; }
   return inverseSqrt(f32({{ ATTN_SCALE_DIM }}));
 }
 {% macro varlen_sequence_of_token(buffer="cumulative_sequence_length") %}
 fn sequence_of_token(token: u32, batch_size: u32) -> u32 {
   for (var b: u32 = 0u; b < batch_size; b = b + 1u) {
@@ -213,23 +156,17 @@ fn sequence_of_token(token: u32, batch_size: u32) -> u32 {
     }
   }
   return batch_size - 1u;
-}
-{%- endmacro %}
 {% macro paged_sequence_index(withLookup=true) %}
-{% if withLookup %}
 {{ varlen_sequence_of_token() }}
-{% endif %}
 // Physical element base of one (slot, kv head) row of the paged cache, whose
 // layout is (num_blocks, block_size, kv_num_heads, head_size).
 fn paged_row_base(block_id: u32, slot_in_block: u32, kv_head: u32) -> u32 {
   return ((block_id * BLOCK_SIZE + slot_in_block) * KV_HEADS + kv_head) * HEAD_DIM;
 }
-{%- endmacro %}
 {{ paged_sequence_index() }}
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
   let token = wg.x;
@@ -403,4 +340,4 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid:
 {% endfor %}
   }
 }
-{%- endif %}

 {% set splitK = splitK is defined and splitK %}
 {% set keyLo = "splitStart" if splitK else "0u" %}
 {% set keyHi = "splitEnd" if splitK else "key_bound" %}
 {{ env.wgsl.resourceDeclarations }}
 // com.microsoft.PagedAttention, attention pass.
   let equalFiniteMax = select(false, value == maxValue, is_finite_f32(maxValue));
   return select(value - maxValue, 0.0, equalFiniteMax);
 }
 fn exp_shift(value: f32, maxValue: f32) -> f32 {
   return exp(shifted_value(value, maxValue));
 }
 // Both the subgroup and portable barrier-tree engines return the same merged
 // pair to every invocation. Repeated merges require a workgroup barrier between
 // calls before their shared partial storage is reused.
 {% set mdStreamed = mdStreams is defined %}
 {% set mdStreams = mdStreams if mdStreams is defined else 1 %}
 {% set mdExtent = "WG" if mdStreams == 1 else "WG * " ~ mdStreams ~ "u" %}
   return merged;
 }
 {% endif %}
+{% set ATTN_SCALE_DIM = "HEAD_DIM" %}fn scale_value() -> f32 {
   if (params.scale != 0.0) { return params.scale; }
   return inverseSqrt(f32({{ ATTN_SCALE_DIM }}));
 }
 {% macro varlen_sequence_of_token(buffer="cumulative_sequence_length") %}
 fn sequence_of_token(token: u32, batch_size: u32) -> u32 {
   for (var b: u32 = 0u; b < batch_size; b = b + 1u) {
     }
   }
   return batch_size - 1u;
+}{% endmacro %}
 {% macro paged_sequence_index(withLookup=true) %}
 {{ varlen_sequence_of_token() }}
 // Physical element base of one (slot, kv head) row of the paged cache, whose
 // layout is (num_blocks, block_size, kv_num_heads, head_size).
 fn paged_row_base(block_id: u32, slot_in_block: u32, kv_head: u32) -> u32 {
   return ((block_id * BLOCK_SIZE + slot_in_block) * KV_HEADS + kv_head) * HEAD_DIM;
 }
+{% endmacro %}
 {{ paged_sequence_index() }}
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
   let token = wg.x;
 {% endfor %}
   }
 }
+{% endif %}

build/webgpu/paged-scatter-kv.wgsl.jinja CHANGED Viewed

@@ -1,37 +1,11 @@
-{% macro flat_index_2d(name="i", bound="params.count", guardInline=false, note="dispatch-limit") %}
-{% if note == "dispatch-limit" %}
-  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
-  // per-axis workgroup fold width (outputs > 16.7M elements).
-{% elif note == "limit" %}
-  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
-  // per-axis workgroup fold width.
-{% elif note == "device-axis" %}
-  // The flat dispatch is folded across x/y at a fixed per-axis workgroup
-  // width; gid.y carries the high portion of the output index.
-{% elif note == "vec4-limit" %}
-  // 2D-folded flat vec4 index: gid.y carries the high bits past the dispatch's
-  // per-axis workgroup fold width (the dispatch caps x and spills into y).
-{% elif note == "element-limit" %}
-  // 2D-folded flat element index: gid.y carries the high bits past the
-  // dispatch's per-axis workgroup fold width.
-{% elif note == "dispatch" %}
   // 2D-folded flat index: gid.y carries the high bits past the dispatch's
   // per-axis workgroup fold width.
-{% endif %}
-{% if bound == "" %}
-  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
-{%- elif guardInline %}
-  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
-  if ({{ name }} >= {{ bound }}) { return; }
-{%- else %}
-  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
   if ({{ name }} >= {{ bound }}) {
     return;
-  }
-{%- endif %}
-{% endmacro %}
-enable f16;
 {{ env.wgsl.resourceDeclarations }}
 // com.microsoft.PagedAttention, KV scatter pass.
@@ -62,26 +36,22 @@ fn sequence_of_token(token: u32, batch_size: u32) -> u32 {
     }
   }
   return batch_size - 1u;
-}
-{%- endmacro %}
 {% macro paged_sequence_index(withLookup=true) %}
 {% if withLookup %}
 {{ varlen_sequence_of_token() }}
-{% endif %}
 // Physical element base of one (slot, kv head) row of the paged cache, whose
 // layout is (num_blocks, block_size, kv_num_heads, head_size).
 fn paged_row_base(block_id: u32, slot_in_block: u32, kv_head: u32) -> u32 {
   return ((block_id * BLOCK_SIZE + slot_in_block) * KV_HEADS + kv_head) * HEAD_DIM;
 }
-{%- endmacro %}
 {{ paged_sequence_index(withLookup=not hasSlotMapping) }}
 @compute @workgroup_size({{ tunables.SCATTER_WORKGROUP_SIZE }})
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-{{ flat_index_2d(bound="params.scatterCount", note="element-limit") }}
   let dim = i % HEAD_DIM;
   let kv_head = (i / HEAD_DIM) % KV_HEADS;
   let token = i / (KV_HEADS * HEAD_DIM);

+{% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
+{% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
   // 2D-folded flat index: gid.y carries the high bits past the dispatch's
   // per-axis workgroup fold width.
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
   if ({{ name }} >= {{ bound }}) {
     return;
+  }{% endmacro %}
 {{ env.wgsl.resourceDeclarations }}
 // com.microsoft.PagedAttention, KV scatter pass.
     }
   }
   return batch_size - 1u;
+}{% endmacro %}
 {% macro paged_sequence_index(withLookup=true) %}
 {% if withLookup %}
 {{ varlen_sequence_of_token() }}
+{% endif %}
 // Physical element base of one (slot, kv head) row of the paged cache, whose
 // layout is (num_blocks, block_size, kv_num_heads, head_size).
 fn paged_row_base(block_id: u32, slot_in_block: u32, kv_head: u32) -> u32 {
   return ((block_id * BLOCK_SIZE + slot_in_block) * KV_HEADS + kv_head) * HEAD_DIM;
 }
+{% endmacro %}
 {{ paged_sequence_index(withLookup=not hasSlotMapping) }}
 @compute @workgroup_size({{ tunables.SCATTER_WORKGROUP_SIZE }})
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+{{ flat_index_2d(tunables.SCATTER_WORKGROUP_SIZE, bound="params.scatterCount") }}
   let dim = i % HEAD_DIM;
   let kv_head = (i / HEAD_DIM) % KV_HEADS;
   let token = i / (KV_HEADS * HEAD_DIM);

build/webgpu/test.json CHANGED Viewed

@@ -686,7 +686,7 @@
     {
       "name": "splitk_slot_mapping",
       "provenance": {
-        "notes": "A 201-key history over a 256-key table creates two split-K ranges, while explicit slot mapping selects the corresponding cache-scatter path."
       },
       "attrs": { "num_heads": 2, "kv_num_heads": 1 },
       "inputs": {
@@ -842,6 +842,341 @@
         }
       },
       "outputs": { "outputT": { "dtype": "float16", "shape": [1, 2048], "tolerance": 0.00001, "relTolerance": 0.0001 } }
     }
   ]
 }

     {
       "name": "splitk_slot_mapping",
       "provenance": {
+        "notes": "A 201-key history over a 256-key cache table with explicit slot mapping checks attention over multiple key ranges and cache publication."
       },
       "attrs": { "num_heads": 2, "kv_num_heads": 1 },
       "inputs": {
         }
       },
       "outputs": { "outputT": { "dtype": "float16", "shape": [1, 2048], "tolerance": 0.00001, "relTolerance": 0.0001 } }
+    },
+    {
+      "name": "scatter_fold_wg1_independent_sequences",
+      "provenance": {
+        "notes": "A large scatter workload crosses the WebGPU dispatch-grid limit. One-token sequences bound attention cost; nonzero dyadic cycles expose missing tail writes and misplaced cache/value elements."
+      },
+      "attrs": { "num_heads": 8, "kv_num_heads": 8 },
+      "tunables": { "WORKGROUP_SIZE": 64, "SCATTER_WORKGROUP_SIZE": 1 },
+      "inputs": {
+        "queryT": { "dtype": "float16", "shape": [64, 1024], "data": { "kind": "constant", "value": 0.0 } },
+        "keyT": {
+          "dtype": "float16",
+          "shape": [64, 1024],
+          "data": { "kind": "cycle", "values": [1.0, 2.0, 4.0, -1.0, -2.0] }
+        },
+        "valueT": {
+          "dtype": "float16",
+          "shape": [64, 1024],
+          "data": { "kind": "cycle", "values": [2.0, -4.0, 8.0, 1.0, -2.0, 4.0, -1.0] }
+        },
+        "keyCacheT": { "dtype": "float16", "shape": [64, 1, 8, 128], "data": { "kind": "constant", "value": 0.0 } },
+        "valueCacheT": { "dtype": "float16", "shape": [64, 1, 8, 128], "data": { "kind": "constant", "value": 0.0 } },
+        "cumulativeSequenceLengthT": {
+          "dtype": "int32",
+          "shape": [65],
+          "data": { "kind": "linspace", "start": 0, "end": 64 }
+        },
+        "pastSeqlensT": { "dtype": "int32", "shape": [64], "data": { "kind": "constant", "value": 0 } },
+        "blockTableT": { "dtype": "int32", "shape": [64, 1], "data": { "kind": "linspace", "start": 0, "end": 63 } }
+      },
+      "outputs": {
+        "outputT": {
+          "dtype": "float16",
+          "shape": [64, 1024],
+          "data": { "kind": "cycle", "values": [2.0, -4.0, 8.0, 1.0, -2.0, 4.0, -1.0] },
+          "tolerance": 0,
+          "relTolerance": 0
+        },
+        "keyCacheT": {
+          "dtype": "float16",
+          "shape": [64, 1, 8, 128],
+          "data": { "kind": "cycle", "values": [1.0, 2.0, 4.0, -1.0, -2.0] },
+          "tolerance": 0,
+          "relTolerance": 0
+        },
+        "valueCacheT": {
+          "dtype": "float16",
+          "shape": [64, 1, 8, 128],
+          "data": { "kind": "cycle", "values": [2.0, -4.0, 8.0, 1.0, -2.0, 4.0, -1.0] },
+          "tolerance": 0,
+          "relTolerance": 0
+        }
+      }
+    },
+    {
+      "name": "scatter_fold_wg4_independent_sequences",
+      "provenance": {
+        "notes": "A large scatter workload crosses the WebGPU dispatch-grid limit. One-token sequences bound attention cost; nonzero dyadic cycles expose missing tail writes and misplaced cache/value elements."
+      },
+      "attrs": { "num_heads": 8, "kv_num_heads": 8 },
+      "tunables": { "WORKGROUP_SIZE": 64, "SCATTER_WORKGROUP_SIZE": 4 },
+      "inputs": {
+        "queryT": { "dtype": "float16", "shape": [256, 1024], "data": { "kind": "constant", "value": 0.0 } },
+        "keyT": {
+          "dtype": "float16",
+          "shape": [256, 1024],
+          "data": { "kind": "cycle", "values": [1.0, 2.0, 4.0, -1.0, -2.0] }
+        },
+        "valueT": {
+          "dtype": "float16",
+          "shape": [256, 1024],
+          "data": { "kind": "cycle", "values": [2.0, -4.0, 8.0, 1.0, -2.0, 4.0, -1.0] }
+        },
+        "keyCacheT": { "dtype": "float16", "shape": [256, 1, 8, 128], "data": { "kind": "constant", "value": 0.0 } },
+        "valueCacheT": { "dtype": "float16", "shape": [256, 1, 8, 128], "data": { "kind": "constant", "value": 0.0 } },
+        "cumulativeSequenceLengthT": {
+          "dtype": "int32",
+          "shape": [257],
+          "data": { "kind": "linspace", "start": 0, "end": 256 }
+        },
+        "pastSeqlensT": { "dtype": "int32", "shape": [256], "data": { "kind": "constant", "value": 0 } },
+        "blockTableT": { "dtype": "int32", "shape": [256, 1], "data": { "kind": "linspace", "start": 0, "end": 255 } }
+      },
+      "outputs": {
+        "outputT": {
+          "dtype": "float16",
+          "shape": [256, 1024],
+          "data": { "kind": "cycle", "values": [2.0, -4.0, 8.0, 1.0, -2.0, 4.0, -1.0] },
+          "tolerance": 0,
+          "relTolerance": 0
+        },
+        "keyCacheT": {
+          "dtype": "float16",
+          "shape": [256, 1, 8, 128],
+          "data": { "kind": "cycle", "values": [1.0, 2.0, 4.0, -1.0, -2.0] },
+          "tolerance": 0,
+          "relTolerance": 0
+        },
+        "valueCacheT": {
+          "dtype": "float16",
+          "shape": [256, 1, 8, 128],
+          "data": { "kind": "cycle", "values": [2.0, -4.0, 8.0, 1.0, -2.0, 4.0, -1.0] },
+          "tolerance": 0,
+          "relTolerance": 0
+        }
+      }
+    },
+    {
+      "name": "ort_paged_ragged_empty_request_b2_h8kv4_d128",
+      "attrs": { "num_heads": 8, "kv_num_heads": 4 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float16",
+          "shape": [2, 1024],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.029, "scale": 0.5 }
+        },
+        "keyT": {
+          "dtype": "float16",
+          "shape": [2, 512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.023, "cosStep": 0.031, "scale": 0.5 }
+        },
+        "valueT": {
+          "dtype": "float16",
+          "shape": [2, 512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.037, "scale": 0.5 }
+        },
+        "keyCacheT": {
+          "dtype": "float16",
+          "shape": [4, 8, 4, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.019, "scale": 0.25 }
+        },
+        "valueCacheT": {
+          "dtype": "float16",
+          "shape": [4, 8, 4, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.007, "cosStep": 0.041, "scale": 0.25 }
+        },
+        "cumulativeSequenceLengthT": {
+          "dtype": "int32",
+          "shape": [3],
+          "data": { "kind": "values", "values": [0, 2, 2] }
+        },
+        "pastSeqlensT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [0, 0] } },
+        "blockTableT": { "dtype": "int32", "shape": [2, 2], "data": { "kind": "values", "values": [2, 0, 3, 1] } }
+      },
+      "outputs": { "outputT": { "dtype": "float16", "shape": [2, 1024], "tolerance": 0.005 } }
+    },
+    {
+      "name": "ort_paged_ragged_new_seqlens_1_2_b2_h8kv4_d128",
+      "attrs": { "num_heads": 8, "kv_num_heads": 4 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float16",
+          "shape": [3, 1024],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.027, "scale": 0.5 }
+        },
+        "keyT": {
+          "dtype": "float16",
+          "shape": [3, 512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.021, "cosStep": 0.033, "scale": 0.5 }
+        },
+        "valueT": {
+          "dtype": "float16",
+          "shape": [3, 512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.015, "cosStep": 0.039, "scale": 0.5 }
+        },
+        "keyCacheT": {
+          "dtype": "float16",
+          "shape": [4, 8, 4, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.019, "scale": 0.25 }
+        },
+        "valueCacheT": {
+          "dtype": "float16",
+          "shape": [4, 8, 4, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.007, "cosStep": 0.041, "scale": 0.25 }
+        },
+        "cumulativeSequenceLengthT": {
+          "dtype": "int32",
+          "shape": [3],
+          "data": { "kind": "values", "values": [0, 1, 3] }
+        },
+        "pastSeqlensT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [0, 0] } },
+        "blockTableT": { "dtype": "int32", "shape": [2, 2], "data": { "kind": "values", "values": [2, 0, 3, 1] } }
+      },
+      "outputs": { "outputT": { "dtype": "float16", "shape": [3, 1024], "tolerance": 0.005 } }
+    },
+    {
+      "name": "ort_paged_prefill_multibatch_varlen_b2t48_h2kv2d128",
+      "attrs": { "num_heads": 2, "kv_num_heads": 2 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float16",
+          "shape": [48, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.023, "scale": 0.4 }
+        },
+        "keyT": {
+          "dtype": "float16",
+          "shape": [48, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.029, "scale": 0.4 }
+        },
+        "valueT": {
+          "dtype": "float16",
+          "shape": [48, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.031, "scale": 0.4 }
+        },
+        "keyCacheT": {
+          "dtype": "float16",
+          "shape": [6, 16, 2, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.007, "cosStep": 0.037, "scale": 0.2 }
+        },
+        "valueCacheT": {
+          "dtype": "float16",
+          "shape": [6, 16, 2, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.041, "scale": 0.2 }
+        },
+        "cumulativeSequenceLengthT": {
+          "dtype": "int32",
+          "shape": [3],
+          "data": { "kind": "values", "values": [0, 32, 48] }
+        },
+        "pastSeqlensT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [0, 0] } },
+        "blockTableT": { "dtype": "int32", "shape": [2, 2], "data": { "kind": "values", "values": [4, 1, 5, 0] } }
+      },
+      "outputs": { "outputT": { "dtype": "float16", "shape": [48, 128], "tolerance": 0.005 } }
+    },
+    {
+      "name": "ort_paged_empty_tokens_copy_caches",
+      "attrs": { "num_heads": 1, "kv_num_heads": 1 },
+      "inputs": {
+        "queryT": { "dtype": "float16", "shape": [0, 8], "data": { "kind": "values", "values": [] } },
+        "keyT": { "dtype": "float16", "shape": [0, 8], "data": { "kind": "values", "values": [] } },
+        "valueT": { "dtype": "float16", "shape": [0, 8], "data": { "kind": "values", "values": [] } },
+        "keyCacheT": {
+          "dtype": "float16",
+          "shape": [2, 4, 1, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.09, "cosStep": 0.13, "scale": 0.5 }
+        },
+        "valueCacheT": {
+          "dtype": "float16",
+          "shape": [2, 4, 1, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.17, "scale": 0.5 }
+        },
+        "cumulativeSequenceLengthT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [0, 0] } },
+        "pastSeqlensT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [0] } },
+        "blockTableT": { "dtype": "int32", "shape": [1, 1], "data": { "kind": "values", "values": [0] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float16", "shape": [0, 8], "data": { "kind": "values", "values": [] } },
+        "keyCacheT": { "dtype": "float16", "shape": [2, 4, 1, 8], "tolerance": 0.001 },
+        "valueCacheT": { "dtype": "float16", "shape": [2, 4, 1, 8], "tolerance": 0.001 }
+      }
+    },
+    {
+      "name": "ort_paged_mixed_prefill_decode_b2t5_past2_4_h2kv1d8",
+      "attrs": { "num_heads": 2, "kv_num_heads": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float16",
+          "shape": [5, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 }
+        },
+        "keyT": {
+          "dtype": "float16",
+          "shape": [5, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 }
+        },
+        "valueT": {
+          "dtype": "float16",
+          "shape": [5, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.37, "scale": 0.5 }
+        },
+        "keyCacheT": {
+          "dtype": "float16",
+          "shape": [4, 4, 1, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.43, "scale": 0.25 }
+        },
+        "valueCacheT": {
+          "dtype": "float16",
+          "shape": [4, 4, 1, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.05, "cosStep": 0.47, "scale": 0.25 }
+        },
+        "cumulativeSequenceLengthT": {
+          "dtype": "int32",
+          "shape": [3],
+          "data": { "kind": "values", "values": [0, 3, 5] }
+        },
+        "pastSeqlensT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [2, 4] } },
+        "blockTableT": { "dtype": "int32", "shape": [2, 2], "data": { "kind": "values", "values": [3, 1, 0, 2] } }
+      },
+      "outputs": { "outputT": { "dtype": "float16", "shape": [5, 16], "tolerance": 0.005 } }
+    },
+    {
+      "name": "ort_paged_splitk_decode_ragged_batch_b2_past200_400_h2kv1d64",
+      "attrs": { "num_heads": 2, "kv_num_heads": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float16",
+          "shape": [2, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.023, "scale": 0.4 }
+        },
+        "keyT": {
+          "dtype": "float16",
+          "shape": [2, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.029, "scale": 0.4 }
+        },
+        "valueT": {
+          "dtype": "float16",
+          "shape": [2, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.031, "scale": 0.4 }
+        },
+        "keyCacheT": {
+          "dtype": "float16",
+          "shape": [32, 32, 1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.003, "cosStep": 0.047, "scale": 0.2 }
+        },
+        "valueCacheT": {
+          "dtype": "float16",
+          "shape": [32, 32, 1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.005, "cosStep": 0.053, "scale": 0.2 }
+        },
+        "cumulativeSequenceLengthT": {
+          "dtype": "int32",
+          "shape": [3],
+          "data": { "kind": "values", "values": [0, 1, 2] }
+        },
+        "pastSeqlensT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [200, 400] } },
+        "blockTableT": {
+          "dtype": "int32",
+          "shape": [2, 16],
+          "data": {
+            "kind": "values",
+            "values": [3, 10, 17, 24, 31, 6, 13, 20, 27, 2, 9, 16, 23, 30, 5, 12, 1, 8, 15, 22, 29, 4, 11, 18, 25, 0, 7, 14, 21, 28, 19, 26]
+          }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float16", "shape": [2, 128], "tolerance": 0.00043 } }
     }
   ]
 }