From 62abea66fa059655ee84e9ba904d45ed96923e13 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 5 Aug 2026 12:54:54 -0700 Subject: [PATCH 1/4] Update (base update) [ghstack-poisoned] --- .../serialization/vulkan_graph_builder.py | 48 +- backends/vulkan/test/test_serialization.py | 138 +- .../test/test_vulkan_compile_options.py | 25 +- backends/vulkan/vulkan_preprocess.py | 5 + backends/webgpu/CMakeLists.txt | 70 +- backends/webgpu/runtime/WebGPUBackend.cpp | 90 +- backends/webgpu/runtime/WebGPUBackend.h | 5 + .../webgpu/runtime/WebGPUDelegateHeader.cpp | 17 +- .../webgpu/runtime/WebGPUDelegateHeader.h | 5 +- backends/webgpu/runtime/WebGPUDispatchMath.h | 128 ++ .../webgpu/runtime/WebGPUExecutionOptions.cpp | 143 ++ .../webgpu/runtime/WebGPUExecutionOptions.h | 87 + backends/webgpu/runtime/WebGPUGraph.cpp | 1939 ++++++++--------- backends/webgpu/runtime/WebGPUGraph.h | 285 ++- .../webgpu/runtime/WebGPUShaderRegistry.cpp | 1108 ++++++++++ .../webgpu/runtime/WebGPUShaderRegistry.h | 26 + backends/webgpu/runtime/WebGPUUtils.h | 66 +- backends/webgpu/runtime/ops/amax/Reduce.cpp | 2 +- backends/webgpu/runtime/ops/amin/Reduce.cpp | 2 +- backends/webgpu/runtime/ops/amin/amin.wgsl | 49 - .../binary_floor_divide_wgsl.h | 2 +- .../binary_minimum_wgsl.h | 2 +- .../ops/{mul => binary_op}/binary_mul_wgsl.h | 2 +- .../runtime/ops/binary_op/binary_op.wgsl | 25 +- .../runtime/ops/binary_op/binary_op.yaml | 23 + .../ops/{pow => binary_op}/binary_pow_wgsl.h | 2 +- .../webgpu/runtime/ops/compare/Compare.cpp | 19 +- .../webgpu/runtime/ops/compare/compare.wgsl | 30 +- .../webgpu/runtime/ops/compare/compare_wgsl.h | 32 +- .../ops/embedding_q4gsw/EmbeddingQ4gsw.cpp | 58 +- .../runtime/ops/{amax => extrema}/amax_wgsl.h | 2 +- .../runtime/ops/{amin => extrema}/amin_wgsl.h | 2 +- .../{amax/amax.wgsl => extrema/extrema.wgsl} | 4 +- .../webgpu/runtime/ops/extrema/extrema.yaml | 13 + .../runtime/ops/floor_divide/BinaryOp.cpp | 2 +- .../ops/floor_divide/binary_floor_divide.wgsl | 51 - backends/webgpu/runtime/ops/linear/Linear.cpp | 18 +- .../runtime/ops/logical_and/LogicalAnd.cpp | 2 +- .../runtime/ops/logical_and/logical_and.wgsl | 25 - .../logical_and_wgsl.h | 2 +- .../logical_binary.wgsl} | 7 +- .../ops/logical_binary/logical_binary.yaml | 13 + .../logical_or_wgsl.h | 2 +- .../runtime/ops/logical_or/LogicalOr.cpp | 2 +- .../webgpu/runtime/ops/minimum/BinaryOp.cpp | 2 +- .../runtime/ops/minimum/binary_minimum.wgsl | 51 - backends/webgpu/runtime/ops/mul/BinaryOp.cpp | 2 +- .../webgpu/runtime/ops/mul/binary_mul.wgsl | 51 - .../runtime/ops/mul/silu_mul_fused.wgsl | 15 +- .../runtime/ops/mul/silu_mul_fused_wgsl.h | 17 +- backends/webgpu/runtime/ops/pow/BinaryOp.cpp | 2 +- .../webgpu/runtime/ops/pow/binary_pow.wgsl | 51 - .../ops/quantized_linear/QuantizedLinear.cpp | 551 +++-- .../q4gsw_linear_gemm_steel.wgsl | 103 +- .../q4gsw_linear_gemm_steel.yaml | 6 + ..._linear_gemm_steel_half_pwdq_f16acc_wgsl.h | 5 +- .../q4gsw_linear_gemm_steel_half_pwdq_wgsl.h | 5 +- .../q4gsw_linear_gemm_steel_half_wgsl.h | 3 +- .../q4gsw_linear_gemm_steel_wgsl.h | 3 +- ...emm_qkv_fused.wgsl => q4gsw_qkv_bk64.wgsl} | 53 +- ...qkv_fused_wgsl.h => q4gsw_qkv_bk64_wgsl.h} | 65 +- .../quantized_linear/q4gsw_steel_bk64_wgsl.h | 158 ++ .../runtime/ops/rope/RotaryEmbedding.cpp | 953 ++++---- .../runtime/ops/rope/rotary_embedding_hf.wgsl | 24 +- .../ops/rope/rotary_embedding_hf_wgsl.h | 26 +- backends/webgpu/runtime/ops/sdpa/Sdpa.cpp | 961 +++++--- .../streaming_attention_k16_causal_bound.wgsl | 268 +++ ...treaming_attention_k16_causal_bound_wgsl.h | 292 +++ ...ming_attention_qwen3_k16_causal_bound.wgsl | 283 +++ ...ng_attention_qwen3_k16_causal_bound_wgsl.h | 311 +++ ..._attention_qwen3_q32_k16_causal_bound.wgsl | 283 +++ ...ttention_qwen3_q32_k16_causal_bound_wgsl.h | 311 +++ .../ops/sdpa_fd_decode/SdpaFdDecode.cpp | 166 +- .../runtime/ops/sdpa_fd_decode/SdpaFdDecode.h | 46 +- .../webgpu/runtime/ops/to_copy/ToCopy.cpp | 99 +- .../ops/to_copy/to_copy_bool_to_float.wgsl | 24 + .../ops/to_copy/to_copy_bool_to_float_wgsl.h | 48 + ...float_to_int.wgsl => to_copy_convert.wgsl} | 6 +- .../runtime/ops/to_copy/to_copy_convert.yaml | 15 + .../ops/to_copy/to_copy_float_to_int_wgsl.h | 2 +- .../ops/to_copy/to_copy_int_to_float.wgsl | 18 - .../ops/to_copy/to_copy_int_to_float_wgsl.h | 2 +- backends/webgpu/runtime/ops/unary/abs.wgsl | 21 - backends/webgpu/runtime/ops/unary/abs_wgsl.h | 2 +- backends/webgpu/runtime/ops/unary/cos_wgsl.h | 2 +- backends/webgpu/runtime/ops/unary/exp.wgsl | 21 - backends/webgpu/runtime/ops/unary/exp_wgsl.h | 2 +- .../webgpu/runtime/ops/unary/hardswish.wgsl | 21 - .../webgpu/runtime/ops/unary/hardswish_wgsl.h | 2 +- backends/webgpu/runtime/ops/unary/neg.wgsl | 21 - backends/webgpu/runtime/ops/unary/neg_wgsl.h | 2 +- backends/webgpu/runtime/ops/unary/round.wgsl | 21 - .../webgpu/runtime/ops/unary/round_wgsl.h | 2 +- backends/webgpu/runtime/ops/unary/rsqrt.wgsl | 21 - .../webgpu/runtime/ops/unary/rsqrt_wgsl.h | 2 +- backends/webgpu/runtime/ops/unary/sin.wgsl | 21 - backends/webgpu/runtime/ops/unary/sin_wgsl.h | 2 +- backends/webgpu/runtime/ops/unary/sqrt.wgsl | 21 - backends/webgpu/runtime/ops/unary/sqrt_wgsl.h | 2 +- backends/webgpu/runtime/ops/unary/tanh.wgsl | 21 - backends/webgpu/runtime/ops/unary/tanh_wgsl.h | 2 +- .../ops/unary/{cos.wgsl => unary.wgsl} | 2 +- backends/webgpu/runtime/ops/unary/unary.yaml | 29 + backends/webgpu/runtime/passes/QkvBk64.cpp | 444 ++++ backends/webgpu/runtime/passes/QkvBk64.h | 71 + backends/webgpu/runtime/passes/SwiGLU.cpp | 338 +++ backends/webgpu/runtime/passes/SwiGLU.h | 59 + backends/webgpu/scripts/gen_wgsl_headers.py | 349 ++- .../webgpu/scripts/test_webgpu_native_ci.sh | 210 +- .../test/native/test_compute_dispatch.cpp | 947 ++++++++ .../webgpu/test/native/test_dispatch_2d.cpp | 261 +++ .../webgpu/test/native/test_dynamic_shape.cpp | 1535 ++++++++++++- .../test/native/test_execution_options.cpp | 167 ++ .../test/native/test_output_suppression.cpp | 180 ++ backends/webgpu/test/op_tests/cases.py | 263 ++- .../webgpu/test/op_tests/generate_op_tests.py | 7 +- .../webgpu/test/op_tests/op_test_driver.cpp | 16 +- .../webgpu/test/op_tests/test_generator.py | 115 +- .../test_dynamic_shape_export.py | 1316 ++++++++++- backends/webgpu/test/ops/test_bitwise.py | 20 +- backends/webgpu/test/ops/test_floor_divide.py | 4 +- backends/webgpu/test/ops/test_logical_and.py | 51 +- backends/webgpu/test/ops/test_logical_or.py | 42 +- backends/webgpu/test/ops/test_minimum.py | 4 +- backends/webgpu/test/ops/test_pow.py | 4 +- .../webgpu/test/ops/test_quantized_linear.py | 221 +- backends/webgpu/test/ops/test_reduce.py | 108 +- backends/webgpu/test/ops/test_rope_hf.py | 282 ++- backends/webgpu/test/ops/test_sdpa.py | 130 +- backends/webgpu/test/ops/test_to_copy.py | 141 +- backends/webgpu/test/test_build_webgpu.sh | 12 + .../webgpu/test/test_cmake_configuration.py | 87 + .../webgpu/test/test_native_ci_contract.py | 75 + backends/webgpu/test/test_webgpu_native.cpp | 1148 +++++++++- backends/webgpu/test/test_wgsl_codegen.py | 707 +++++- 135 files changed, 15969 insertions(+), 3444 deletions(-) create mode 100644 backends/webgpu/runtime/WebGPUExecutionOptions.cpp create mode 100644 backends/webgpu/runtime/WebGPUExecutionOptions.h create mode 100644 backends/webgpu/runtime/WebGPUShaderRegistry.cpp create mode 100644 backends/webgpu/runtime/WebGPUShaderRegistry.h delete mode 100644 backends/webgpu/runtime/ops/amin/amin.wgsl rename backends/webgpu/runtime/ops/{floor_divide => binary_op}/binary_floor_divide_wgsl.h (97%) rename backends/webgpu/runtime/ops/{minimum => binary_op}/binary_minimum_wgsl.h (97%) rename backends/webgpu/runtime/ops/{mul => binary_op}/binary_mul_wgsl.h (98%) rename backends/webgpu/runtime/ops/{pow => binary_op}/binary_pow_wgsl.h (98%) rename backends/webgpu/runtime/ops/{amax => extrema}/amax_wgsl.h (97%) rename backends/webgpu/runtime/ops/{amin => extrema}/amin_wgsl.h (97%) rename backends/webgpu/runtime/ops/{amax/amax.wgsl => extrema/extrema.wgsl} (94%) create mode 100644 backends/webgpu/runtime/ops/extrema/extrema.yaml delete mode 100644 backends/webgpu/runtime/ops/floor_divide/binary_floor_divide.wgsl delete mode 100644 backends/webgpu/runtime/ops/logical_and/logical_and.wgsl rename backends/webgpu/runtime/ops/{logical_and => logical_binary}/logical_and_wgsl.h (96%) rename backends/webgpu/runtime/ops/{logical_or/logical_or.wgsl => logical_binary/logical_binary.wgsl} (73%) create mode 100644 backends/webgpu/runtime/ops/logical_binary/logical_binary.yaml rename backends/webgpu/runtime/ops/{logical_or => logical_binary}/logical_or_wgsl.h (96%) delete mode 100644 backends/webgpu/runtime/ops/minimum/binary_minimum.wgsl delete mode 100644 backends/webgpu/runtime/ops/mul/binary_mul.wgsl delete mode 100644 backends/webgpu/runtime/ops/pow/binary_pow.wgsl rename backends/webgpu/runtime/ops/quantized_linear/{q4gsw_linear_gemm_qkv_fused.wgsl => q4gsw_qkv_bk64.wgsl} (63%) rename backends/webgpu/runtime/ops/quantized_linear/{q4gsw_linear_gemm_qkv_fused_wgsl.h => q4gsw_qkv_bk64_wgsl.h} (61%) create mode 100644 backends/webgpu/runtime/ops/quantized_linear/q4gsw_steel_bk64_wgsl.h create mode 100644 backends/webgpu/runtime/ops/sdpa/streaming_attention_k16_causal_bound.wgsl create mode 100644 backends/webgpu/runtime/ops/sdpa/streaming_attention_k16_causal_bound_wgsl.h create mode 100644 backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_k16_causal_bound.wgsl create mode 100644 backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_k16_causal_bound_wgsl.h create mode 100644 backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_q32_k16_causal_bound.wgsl create mode 100644 backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_q32_k16_causal_bound_wgsl.h create mode 100644 backends/webgpu/runtime/ops/to_copy/to_copy_bool_to_float.wgsl create mode 100644 backends/webgpu/runtime/ops/to_copy/to_copy_bool_to_float_wgsl.h rename backends/webgpu/runtime/ops/to_copy/{to_copy_float_to_int.wgsl => to_copy_convert.wgsl} (62%) create mode 100644 backends/webgpu/runtime/ops/to_copy/to_copy_convert.yaml delete mode 100644 backends/webgpu/runtime/ops/to_copy/to_copy_int_to_float.wgsl delete mode 100644 backends/webgpu/runtime/ops/unary/abs.wgsl delete mode 100644 backends/webgpu/runtime/ops/unary/exp.wgsl delete mode 100644 backends/webgpu/runtime/ops/unary/hardswish.wgsl delete mode 100644 backends/webgpu/runtime/ops/unary/neg.wgsl delete mode 100644 backends/webgpu/runtime/ops/unary/round.wgsl delete mode 100644 backends/webgpu/runtime/ops/unary/rsqrt.wgsl delete mode 100644 backends/webgpu/runtime/ops/unary/sin.wgsl delete mode 100644 backends/webgpu/runtime/ops/unary/sqrt.wgsl delete mode 100644 backends/webgpu/runtime/ops/unary/tanh.wgsl rename backends/webgpu/runtime/ops/unary/{cos.wgsl => unary.wgsl} (94%) create mode 100644 backends/webgpu/runtime/ops/unary/unary.yaml create mode 100644 backends/webgpu/runtime/passes/QkvBk64.cpp create mode 100644 backends/webgpu/runtime/passes/QkvBk64.h create mode 100644 backends/webgpu/runtime/passes/SwiGLU.cpp create mode 100644 backends/webgpu/runtime/passes/SwiGLU.h create mode 100644 backends/webgpu/test/native/test_compute_dispatch.cpp create mode 100644 backends/webgpu/test/native/test_execution_options.cpp create mode 100644 backends/webgpu/test/native/test_output_suppression.cpp create mode 100644 backends/webgpu/test/test_cmake_configuration.py create mode 100644 backends/webgpu/test/test_native_ci_contract.py diff --git a/backends/vulkan/serialization/vulkan_graph_builder.py b/backends/vulkan/serialization/vulkan_graph_builder.py index ca5ab196dd2..46e01e701b1 100644 --- a/backends/vulkan/serialization/vulkan_graph_builder.py +++ b/backends/vulkan/serialization/vulkan_graph_builder.py @@ -9,7 +9,7 @@ import logging import operator from types import NoneType -from typing import cast, List, Optional, Union +from typing import cast, Dict, List, Optional, Union import executorch.backends.vulkan.serialization.vulkan_graph_schema as vk_graph_schema import torch @@ -28,6 +28,7 @@ ) from executorch.exir._serialize._named_data_store import NamedDataStore from executorch.exir.backend.utils import DelegateMappingBuilder +from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.tensor import TensorSpec from torch._export.utils import get_buffer, get_param, is_buffer, is_param from torch.export import ExportedProgram @@ -49,11 +50,41 @@ def __init__( delegate_mapping_builder: DelegateMappingBuilder, downcast_64_bit: bool = True, force_fp16: bool = False, + alias_buffer_mutations: bool = False, ) -> None: self.program = program self.delegate_mapping_builder = delegate_mapping_builder self.downcast_64_bit = downcast_64_bit self.force_fp16 = force_fp16 + self.buffer_mutation_inputs: Dict[str, Node] = {} + self.buffer_mutation_user_outputs: set[str] = set() + if alias_buffer_mutations: + nodes_by_name = { + node.name: node for node in program.graph_module.graph.nodes + } + buffer_inputs_by_target: Dict[str, Node] = {} + for name, target in program.graph_signature.inputs_to_buffers.items(): + if name not in nodes_by_name: + continue + buffer_input = nodes_by_name[name] + prepack = next( + ( + user + for user in buffer_input.users + if user.op == "call_function" + and user.target == exir_ops.edge.et_vk.prepack.default + ), + None, + ) + buffer_inputs_by_target[target] = prepack or buffer_input + self.buffer_mutation_inputs = { + output_name: buffer_inputs_by_target[target] + for output_name, target in program.graph_signature.buffers_to_mutate.items() + if target in buffer_inputs_by_target + } + self.buffer_mutation_user_outputs = set( + program.graph_signature.user_outputs + ) self.chain = [] self.values = [] self.input_ids = [] @@ -160,6 +191,16 @@ def maybe_add_constant_tensor(self, node: Node) -> int: return constant_id def create_node_value(self, node: Node) -> int: + if node.name in self.buffer_mutation_inputs: + input_node = self.buffer_mutation_inputs[node.name] + if input_node not in self.node_to_value_ids: + raise AssertionError( + "Cannot alias a buffer mutation before its input is serialized" + ) + value_id = self.node_to_value_ids[input_node] + self.node_to_value_ids[node] = value_id + return value_id + # If the node has been marked as a scalar tensor, create a SymInt instead of a tensor if is_symint_node(node) or node.meta.get("etvk_is_scalar_tensor", False): new_id = self.create_symint_value() @@ -448,7 +489,10 @@ def process_output_node(self, node: Node) -> None: ) # Mutable buffers outputs are not included as an output to the # delegate call. Skip marking them as an output. - if is_mutable_buffer_node(out_node, self.program): + if out_node.name in self.buffer_mutation_inputs: + if out_node.name not in self.buffer_mutation_user_outputs: + continue + elif is_mutable_buffer_node(out_node, self.program): continue self.output_ids.append(self.node_to_value_ids[out_node]) diff --git a/backends/vulkan/test/test_serialization.py b/backends/vulkan/test/test_serialization.py index c373f5216d2..71a6980635a 100644 --- a/backends/vulkan/test/test_serialization.py +++ b/backends/vulkan/test/test_serialization.py @@ -9,9 +9,14 @@ import ctypes import random import unittest -from typing import List +from types import SimpleNamespace +from typing import List, Tuple +import executorch.backends.vulkan.custom_ops_lib # noqa: F401 import torch +from executorch.backends.vulkan.serialization import ( + vulkan_graph_builder as graph_builder_module, +) from executorch.backends.vulkan.serialization.vulkan_graph_schema import ( IntList, @@ -30,6 +35,137 @@ class TestSerialization(unittest.TestCase): + def _build_mutation_program( + self, prepack: bool, shared_user_output: bool = False + ) -> Tuple[SimpleNamespace, torch.fx.Node, torch.fx.Node, torch.fx.Node]: + graph = torch.fx.Graph() + state = graph.placeholder("state") + user_input = graph.placeholder("user_input") + state.meta["spec"] = graph_builder_module.TensorSpec.from_tensor(torch.zeros(4)) + user_input.meta["spec"] = graph_builder_module.TensorSpec.from_tensor( + torch.ones(4) + ) + + state_value = state + if prepack: + state_value = graph.call_function( + graph_builder_module.exir_ops.edge.et_vk.prepack.default, + (state,), + ) + state_value.meta["spec"] = graph_builder_module.TensorSpec.from_tensor( + torch.zeros(4) + ) + + mutation = graph.call_function( + torch.ops.aten.add.Tensor, (state_value, user_input) + ) + mutation.meta["spec"] = graph_builder_module.TensorSpec.from_tensor( + torch.ones(4) + ) + user_output = mutation + if not shared_user_output: + user_output = graph.call_function( + torch.ops.aten.mul.Tensor, (user_input, 2.0) + ) + user_output.meta["spec"] = graph_builder_module.TensorSpec.from_tensor( + torch.ones(4) + ) + graph.output((mutation, user_output)) + + graph_module = torch.fx.GraphModule({}, graph) + signature = SimpleNamespace( + buffers_to_mutate={mutation.name: "state"}, + inputs_to_buffers={state.name: "state"}, + inputs_to_lifted_tensor_constants={}, + inputs_to_parameters={}, + non_persistent_buffers=set(), + user_outputs=(user_output.name,), + ) + program = SimpleNamespace( + constants={}, + graph_module=graph_module, + graph_signature=signature, + state_dict={"state": torch.zeros(4)}, + ) + return program, state_value, mutation, user_output + + def test_alias_buffer_mutations_is_opt_in(self) -> None: + for prepack in (False, True): + with self.subTest(prepack=prepack): + program, state_value, mutation, user_output = ( + self._build_mutation_program(prepack) + ) + + default_builder = graph_builder_module.VkGraphBuilder( + program, + graph_builder_module.DelegateMappingBuilder( + generated_identifiers=True + ), + ) + default_graph = default_builder.build_graph() + self.assertNotEqual( + default_builder.node_to_value_ids[mutation], + default_builder.node_to_value_ids[state_value], + ) + self.assertEqual( + default_graph.output_ids, + [ + default_builder.node_to_value_ids[mutation], + default_builder.node_to_value_ids[user_output], + ], + ) + + explicit_false_builder = graph_builder_module.VkGraphBuilder( + program, + graph_builder_module.DelegateMappingBuilder( + generated_identifiers=True + ), + alias_buffer_mutations=False, + ) + self.assertEqual(default_graph, explicit_false_builder.build_graph()) + + aliasing_builder = graph_builder_module.VkGraphBuilder( + program, + graph_builder_module.DelegateMappingBuilder( + generated_identifiers=True + ), + alias_buffer_mutations=True, + ) + aliasing_graph = aliasing_builder.build_graph() + self.assertEqual( + aliasing_builder.node_to_value_ids[mutation], + aliasing_builder.node_to_value_ids[state_value], + ) + self.assertEqual( + aliasing_graph.output_ids, + [aliasing_builder.node_to_value_ids[user_output]], + ) + + def test_alias_buffer_mutations_preserves_shared_user_output(self) -> None: + for prepack in (False, True): + with self.subTest(prepack=prepack): + program, state_value, mutation, _ = self._build_mutation_program( + prepack, shared_user_output=True + ) + builder = graph_builder_module.VkGraphBuilder( + program, + graph_builder_module.DelegateMappingBuilder( + generated_identifiers=True + ), + alias_buffer_mutations=True, + ) + + graph = builder.build_graph() + + self.assertEqual( + builder.node_to_value_ids[mutation], + builder.node_to_value_ids[state_value], + ) + self.assertEqual( + graph.output_ids, + [builder.node_to_value_ids[mutation]], + ) + def _generate_random_const_tensors(self, num_tensors: int) -> List[torch.Tensor]: """ Helper function to generate `num_tensor` buffers of random sizes and random contents, diff --git a/backends/vulkan/test/test_vulkan_compile_options.py b/backends/vulkan/test/test_vulkan_compile_options.py index f45cfdf12d9..ff507bb3197 100644 --- a/backends/vulkan/test/test_vulkan_compile_options.py +++ b/backends/vulkan/test/test_vulkan_compile_options.py @@ -41,6 +41,10 @@ def test_skip_memory_planning_round_trips(self) -> None: round_tripped = self._round_trip({"skip_memory_planning": True}) self.assertTrue(round_tripped.get("skip_memory_planning")) + def test_alias_buffer_mutations_round_trips(self) -> None: + round_tripped = self._round_trip({"alias_buffer_mutations": True}) + self.assertTrue(round_tripped.get("alias_buffer_mutations")) + def test_force_fp16_round_trips(self) -> None: round_tripped = self._round_trip({"force_fp16": True}) self.assertTrue(round_tripped.get("force_fp16")) @@ -105,15 +109,15 @@ def build_graph(): ), patch( "executorch.backends.vulkan.vulkan_preprocess.VkGraphBuilder", return_value=graph_builder, - ), patch( + ) as graph_builder_factory, patch( "executorch.backends.vulkan.vulkan_preprocess.serialize_vulkan_graph", return_value=b"vk_graph", ): result = VulkanBackend.preprocess(program, parse_compile_options(options)) - return result.data_store_output, externalize_pte_data + return result.data_store_output, externalize_pte_data, graph_builder_factory def test_external_constants_default_keeps_constants_inline(self) -> None: - output, externalize_pte_data = self._preprocess_named_data({}) + output, externalize_pte_data, _ = self._preprocess_named_data({}) self.assertEqual(output.buffers, [b"constant"]) self.assertEqual(output.pte_data, {"constant": DataEntry(0, 16, None)}) @@ -121,7 +125,7 @@ def test_external_constants_default_keeps_constants_inline(self) -> None: externalize_pte_data.assert_not_called() def test_external_constants_option_externalizes_constants(self) -> None: - output, externalize_pte_data = self._preprocess_named_data( + output, externalize_pte_data, _ = self._preprocess_named_data( {"external_constants_max_data_bytes": 16} ) @@ -131,8 +135,21 @@ def test_external_constants_option_externalizes_constants(self) -> None: self.assertEqual(list(next(iter(output.external_data.values()))), ["constant"]) externalize_pte_data.assert_called_once_with(16, "vulkan_constants") + def test_alias_buffer_mutations_reaches_graph_builder(self) -> None: + for options, expected in ( + ({}, False), + ({"alias_buffer_mutations": True}, True), + ): + with self.subTest(options=options): + _, _, graph_builder_factory = self._preprocess_named_data(options) + self.assertIs( + graph_builder_factory.call_args.kwargs["alias_buffer_mutations"], + expected, + ) + def test_unset_options_are_absent(self) -> None: round_tripped = self._round_trip({}) + self.assertNotIn("alias_buffer_mutations", round_tripped) self.assertNotIn("small_texture_limits", round_tripped) self.assertNotIn("skip_memory_planning", round_tripped) self.assertNotIn("external_constants_max_data_bytes", round_tripped) diff --git a/backends/vulkan/vulkan_preprocess.py b/backends/vulkan/vulkan_preprocess.py index d3954240880..f7d6955ce26 100644 --- a/backends/vulkan/vulkan_preprocess.py +++ b/backends/vulkan/vulkan_preprocess.py @@ -130,6 +130,9 @@ def parse_compile_spec(compile_specs: List[CompileSpec]) -> Dict[str, Any]: if spec.key == "skip_memory_planning": options[spec.key] = bool.from_bytes(spec.value, byteorder="little") + if spec.key == "alias_buffer_mutations": + options[spec.key] = bool.from_bytes(spec.value, byteorder="little") + if spec.key == "external_constants_max_data_bytes": options[spec.key] = _parse_external_constants_max_data_bytes(spec.value) @@ -172,6 +175,7 @@ def preprocess( # noqa: C901 ) downcast_64_bit = compile_options.get("downcast_64_bit", True) force_fp16 = compile_options.get("force_fp16", False) + alias_buffer_mutations = compile_options.get("alias_buffer_mutations", False) program = unsafe_remove_auto_functionalized_pass(program) @@ -258,6 +262,7 @@ def preprocess( # noqa: C901 DelegateMappingBuilder(generated_identifiers=True), downcast_64_bit=downcast_64_bit, force_fp16=force_fp16, + alias_buffer_mutations=alias_buffer_mutations, ) vk_graph = graph_builder.build_graph() external_constants_max_data_bytes = compile_options.get( diff --git a/backends/webgpu/CMakeLists.txt b/backends/webgpu/CMakeLists.txt index 3060d52bf03..536348ca69f 100644 --- a/backends/webgpu/CMakeLists.txt +++ b/backends/webgpu/CMakeLists.txt @@ -26,9 +26,16 @@ if(NOT TARGET vulkan_schema) endif() set(WEBGPU_SRCS - runtime/WebGPUBackend.cpp runtime/WebGPUGraph.cpp - runtime/WebGPUDelegateHeader.cpp runtime/WebGPUDevice.cpp - runtime/WebGPUQueryPool.cpp runtime/ops/OperatorRegistry.cpp + runtime/WebGPUBackend.cpp + runtime/WebGPUExecutionOptions.cpp + runtime/WebGPUGraph.cpp + runtime/passes/SwiGLU.cpp + runtime/passes/QkvBk64.cpp + runtime/WebGPUDelegateHeader.cpp + runtime/WebGPUDevice.cpp + runtime/WebGPUQueryPool.cpp + runtime/WebGPUShaderRegistry.cpp + runtime/ops/OperatorRegistry.cpp ) # Op handlers: glob so adding an op needs no CMakeLists edit. CONFIGURE_DEPENDS @@ -57,22 +64,27 @@ target_include_directories( target_link_libraries(webgpu_backend PRIVATE vulkan_schema executorch_core) -# Native WebGPU backend: Dawn (Tint) + SwiftShader; deps script sets Dawn_DIR. -# Native-only: browser/Emscripten builds use the system webgpu.h and never reach -# this find_package (root CMake gates it via EXECUTORCH_BUILD_WEBGPU). -# dawn::webgpu_dawn's link interface references Threads::Threads. -find_package(Threads REQUIRED) -find_package(Dawn REQUIRED) -set(WEBGPU_GPU_LIB dawn::webgpu_dawn) -target_link_libraries(webgpu_backend PUBLIC ${WEBGPU_GPU_LIB}) - -if(APPLE) - target_link_libraries( - webgpu_backend PRIVATE "-framework Metal" "-framework QuartzCore" - "-framework CoreGraphics" "-framework Foundation" - ) +# WASM gets its WebGPU implementation from emdawnwebgpu at executable link time. +# Native builds link Dawn (Tint) and the platform GPU libraries. +if(EMSCRIPTEN) + target_compile_options(webgpu_backend PUBLIC "--use-port=emdawnwebgpu") + # --use-port is also required at link time: the link step pulls in the port's + # headers, JS glue, and libraries for consumers that link webgpu_backend. + target_link_options(webgpu_backend PUBLIC "--use-port=emdawnwebgpu") else() - target_link_libraries(webgpu_backend PRIVATE dl m pthread) + find_package(Threads REQUIRED) + find_package(Dawn REQUIRED) + set(WEBGPU_GPU_LIB dawn::webgpu_dawn) + target_link_libraries(webgpu_backend PUBLIC ${WEBGPU_GPU_LIB}) + + if(APPLE) + target_link_libraries( + webgpu_backend PRIVATE "-framework Metal" "-framework QuartzCore" + "-framework CoreGraphics" "-framework Foundation" + ) + else() + target_link_libraries(webgpu_backend PRIVATE dl m pthread) + endif() endif() target_compile_options(webgpu_backend PRIVATE -fexceptions) @@ -201,5 +213,27 @@ if(EXECUTORCH_BUILD_WEBGPU_TEST) target_link_libraries( webgpu_dispatch_2d_test PRIVATE GTest::gtest GTest::gtest_main ) + add_executable( + webgpu_execution_options_test test/native/test_execution_options.cpp + runtime/WebGPUExecutionOptions.cpp + ) + target_include_directories( + webgpu_execution_options_test + PRIVATE $ + ) + target_link_libraries( + webgpu_execution_options_test PRIVATE GTest::gtest GTest::gtest_main + ) + target_compile_options(webgpu_execution_options_test PRIVATE -fexceptions) + set_property(TARGET webgpu_execution_options_test PROPERTY CXX_STANDARD 17) + add_webgpu_native_test( + webgpu_output_suppression_test test/native/test_output_suppression.cpp + ) + target_link_libraries(webgpu_output_suppression_test PRIVATE GTest::gtest) + + add_webgpu_native_test( + webgpu_compute_dispatch_test test/native/test_compute_dispatch.cpp + ) + target_link_libraries(webgpu_compute_dispatch_test PRIVATE GTest::gtest) endif() endif() diff --git a/backends/webgpu/runtime/WebGPUBackend.cpp b/backends/webgpu/runtime/WebGPUBackend.cpp index e7645514c9b..35d4225bc7c 100644 --- a/backends/webgpu/runtime/WebGPUBackend.cpp +++ b/backends/webgpu/runtime/WebGPUBackend.cpp @@ -8,6 +8,7 @@ #include #include +#include #include #include @@ -17,6 +18,7 @@ #include #include +#include #include #include @@ -42,6 +44,27 @@ using executorch::runtime::resize_tensor; using executorch::runtime::Result; using executorch::runtime::Span; +Result parse_webgpu_graph_config( + ArrayRef compile_specs) { + WebGPUGraphConfig config; + for (const CompileSpec& spec : compile_specs) { + if (spec.key == nullptr || + std::strcmp(spec.key, "webgpu_record_q4gsw_decode_route") != 0) { + continue; + } + if (spec.value.nbytes != sizeof(uint8_t) || spec.value.buffer == nullptr) { + ET_LOG( + Error, + "WebGPU compile option webgpu_record_q4gsw_decode_route must be " + "exactly one byte"); + return Error::DelegateInvalidCompatibility; + } + config.record_q4gsw_decode_route = + *static_cast(spec.value.buffer) != 0; + } + return config; +} + bool WebGPUBackend::is_available() const { return true; } @@ -50,6 +73,13 @@ Result WebGPUBackend::init( BackendInitContext& context, FreeableBuffer* processed, ArrayRef compile_specs) const { + Result parsed_config = + parse_webgpu_graph_config(compile_specs); + if (!parsed_config.ok()) { + return parsed_config.error(); + } + WebGPUGraphConfig config = parsed_config.get(); + // Allocate graph on the runtime allocator WebGPUGraph* graph = context.get_runtime_allocator()->allocateInstance(); @@ -60,7 +90,7 @@ Result WebGPUBackend::init( // Parse header to locate flatbuffer and constant data Result header = - WebGPUDelegateHeader::parse(processed->data()); + WebGPUDelegateHeader::parse(processed->data(), processed->size()); if (!header.ok()) { ET_LOG(Error, "WebGPUDelegateHeader may be corrupt"); return header.error(); @@ -71,6 +101,17 @@ Result WebGPUBackend::init( const uint8_t* flatbuffer_data = buffer_start + header->flatbuffer_offset; const uint8_t* constant_data = buffer_start + header->bytes_offset; + size_t constant_data_size = header->bytes_size; + if (constant_data_size == 0 && processed->size() > header->bytes_offset) { + constant_data_size = processed->size() - header->bytes_offset; + } + + flatbuffers::Verifier verifier(flatbuffer_data, header->flatbuffer_size); + if (!vkgraph::VerifyVkGraphBuffer(verifier)) { + ET_LOG(Error, "WebGPU delegate FlatBuffer verification failed"); + return Error::DelegateInvalidCompatibility; + } + // Verify FlatBuffer identifier if (!vkgraph::VkGraphBufferHasIdentifier(flatbuffer_data)) { ET_LOG( @@ -82,19 +123,23 @@ Result WebGPUBackend::init( // Load-time backend option (BackendOption / LoadBackendOptionsMap), keyed by // the registered backend name; default false. Mirrors the CoreML/XNNPACK // runtime-spec pattern -- no compile flag and no .pte re-export needed. - bool enable_f16_kv_cache = false; { Result spec = context.get_runtime_spec("enable_f16_kv_cache"); if (spec.ok()) { - enable_f16_kv_cache = spec.get(); + config.f16_kv_cache = spec.get(); } } - bool enable_f16_accumulate_gemm = false; { Result spec = context.get_runtime_spec("enable_f16_accumulate_gemm"); if (spec.ok()) { - enable_f16_accumulate_gemm = spec.get(); + config.f16_accumulate_gemm = spec.get(); + } + } + { + Result spec = context.get_runtime_spec("sdpa_query_tile"); + if (spec.ok()) { + config.sdpa_query_tile = spec.get(); } } @@ -102,9 +147,9 @@ Result WebGPUBackend::init( graph->build( flatbuffer_data, constant_data, + constant_data_size, context.get_named_data_map(), - enable_f16_kv_cache, - enable_f16_accumulate_gemm); + config); } catch (const std::exception& e) { ET_LOG(Error, "WebGPU graph build failed: %s", e.what()); graph->~WebGPUGraph(); @@ -121,9 +166,11 @@ Error WebGPUBackend::execute( DelegateHandle* handle, Span args) const { WebGPUGraph* graph = static_cast(handle); + const WebGPUExecutionOptions options = current_webgpu_execution_options(); const size_t num_inputs = graph->input_ids().size(); const size_t num_outputs = graph->output_ids().size(); + WebGPUGraphExecutionOptions graph_options; // Copy inputs from EValue tensors to GPU buffers std::vector inputs; @@ -137,8 +184,13 @@ Error WebGPUBackend::execute( const auto& tensor = args[i]->toTensor(); const bool host_is_int64 = tensor.scalar_type() == executorch::aten::ScalarType::Long; + const bool host_is_fp32 = + tensor.scalar_type() == executorch::aten::ScalarType::Float; inputs.push_back( - {tensor.const_data_ptr(), tensor.nbytes(), host_is_int64}); + {tensor.const_data_ptr(), + tensor.nbytes(), + host_is_int64, + host_is_fp32}); const auto sizes = tensor.sizes(); std::vector new_dims(sizes.begin(), sizes.end()); graph->resize_input(graph->input_ids()[i], new_dims); @@ -159,6 +211,14 @@ Error WebGPUBackend::execute( return Error::Internal; } } + std::vector delegate_outputs; + delegate_outputs.reserve(num_outputs); + for (size_t i = 0; i < num_outputs; i++) { + delegate_outputs.push_back( + args[num_inputs + i]->toTensor().mutable_data_ptr()); + } + graph_options = + resolve_webgpu_graph_execution_options(delegate_outputs, options); } catch (const std::exception& e) { ET_LOG(Error, "WebGPU input/output resize / copy failed: %s", e.what()); return Error::Internal; @@ -167,15 +227,21 @@ Error WebGPUBackend::execute( // Execute + read back; fail loud as a runtime Error so a throw never crosses // the backend boundary. try { - graph->execute(); - std::vector> outputs; + const WebGPUExecutionPlan plan = graph->make_execution_plan(graph_options); + graph->execute(plan); + + // Copy outputs from GPU staging buffers to EValue tensor data pointers + std::vector outputs; outputs.reserve(num_outputs); for (size_t i = 0; i < num_outputs; i++) { const size_t arg_idx = num_inputs + i; auto& tensor = args[arg_idx]->toTensor(); - outputs.emplace_back(tensor.mutable_data_ptr(), tensor.nbytes()); + const bool host_is_fp32 = + tensor.scalar_type() == executorch::aten::ScalarType::Float; + outputs.push_back( + {tensor.mutable_data_ptr(), tensor.nbytes(), host_is_fp32}); } - graph->copy_outputs(outputs); + graph->copy_outputs(outputs, plan); } catch (const std::exception& e) { ET_LOG(Error, "WebGPU execute / output copy failed: %s", e.what()); return Error::Internal; diff --git a/backends/webgpu/runtime/WebGPUBackend.h b/backends/webgpu/runtime/WebGPUBackend.h index 9c20a3d53be..59f7e33994c 100644 --- a/backends/webgpu/runtime/WebGPUBackend.h +++ b/backends/webgpu/runtime/WebGPUBackend.h @@ -8,12 +8,17 @@ #pragma once +#include #include namespace executorch { namespace backends { namespace webgpu { +executorch::runtime::Result parse_webgpu_graph_config( + executorch::runtime::ArrayRef + compile_specs); + class WebGPUBackend final : public ::executorch::runtime::BackendInterface { public: ~WebGPUBackend() override = default; diff --git a/backends/webgpu/runtime/WebGPUDelegateHeader.cpp b/backends/webgpu/runtime/WebGPUDelegateHeader.cpp index d1e8b2110a7..69e6dd70536 100644 --- a/backends/webgpu/runtime/WebGPUDelegateHeader.cpp +++ b/backends/webgpu/runtime/WebGPUDelegateHeader.cpp @@ -65,13 +65,19 @@ bool WebGPUDelegateHeader::is_valid() const { if (flatbuffer_size == 0) { return false; } - if (bytes_offset < flatbuffer_offset + flatbuffer_size) { + if (bytes_offset < flatbuffer_offset || + flatbuffer_size > bytes_offset - flatbuffer_offset) { return false; } return true; } -Result WebGPUDelegateHeader::parse(const void* data) { +Result WebGPUDelegateHeader::parse( + const void* data, + size_t buffer_size) { + if (data == nullptr || buffer_size < kExpectedSize) { + return Error::InvalidArgument; + } const uint8_t* header_data = (const uint8_t*)data; const uint8_t* magic_start = header_data + kMagic.offset; @@ -91,6 +97,13 @@ Result WebGPUDelegateHeader::parse(const void* data) { return Error::InvalidArgument; } + if (header.flatbuffer_offset > buffer_size || + header.flatbuffer_size > buffer_size - header.flatbuffer_offset || + header.bytes_offset > buffer_size || + header.bytes_size > buffer_size - header.bytes_offset) { + return Error::InvalidArgument; + } + return header; } diff --git a/backends/webgpu/runtime/WebGPUDelegateHeader.h b/backends/webgpu/runtime/WebGPUDelegateHeader.h index 6f2f65130c7..14dcfd14a9a 100644 --- a/backends/webgpu/runtime/WebGPUDelegateHeader.h +++ b/backends/webgpu/runtime/WebGPUDelegateHeader.h @@ -8,6 +8,8 @@ #pragma once +#include + #include namespace executorch { @@ -18,7 +20,8 @@ struct WebGPUDelegateHeader { bool is_valid() const; static executorch::runtime::Result parse( - const void* data); + const void* data, + size_t buffer_size); uint32_t header_size; uint32_t flatbuffer_offset; diff --git a/backends/webgpu/runtime/WebGPUDispatchMath.h b/backends/webgpu/runtime/WebGPUDispatchMath.h index 56544c6d496..561bebc6b87 100644 --- a/backends/webgpu/runtime/WebGPUDispatchMath.h +++ b/backends/webgpu/runtime/WebGPUDispatchMath.h @@ -13,6 +13,7 @@ // requires for its device-facing functions). #include +#include #include #include #include @@ -74,6 +75,133 @@ struct DispatchGrid { uint32_t stride_x; }; +struct WgCount { + uint32_t x; + uint32_t y; +}; + +struct DispatchRange { + size_t begin; + size_t end; +}; + +constexpr bool should_record_q4gsw_dual_route( + uint32_t max_m, + bool bicol_eligible, + bool has_dynamic_shapes, + bool record_q4gsw_decode_route) { + return max_m > 1u && bicol_eligible && + (has_dynamic_shapes || record_q4gsw_decode_route); +} + +constexpr bool should_record_sdpa_dual_route( + bool fd_eligible, + bool has_dynamic_sequence, + bool has_dynamic_position) { + return fd_eligible && (has_dynamic_sequence || has_dynamic_position); +} + +constexpr bool is_q4gsw_bk64_eligible( + uint32_t k, + uint32_t n, + uint32_t group_size, + bool has_bias, + bool shader_f16_supported, + uint32_t max_invocations, + uint32_t max_workgroup_storage_bytes) { + constexpr uint32_t kRequiredInvocations = 256u; + constexpr uint32_t kRequiredStorageBytes = 2u * 64u * 64u * sizeof(uint16_t); + const bool ordinary_llama_projection = (k == 2048u && n == 8192u) || + (k == 8192u && n == 2048u) || (k == 2048u && n == 2048u); + return ordinary_llama_projection && k % 64u == 0u && group_size == 64u && + !has_bias && shader_f16_supported && + max_invocations >= kRequiredInvocations && + max_workgroup_storage_bytes >= kRequiredStorageBytes; +} + +constexpr bool is_q4gsw_bk64_live_m(uint32_t m) { + return m == 128u || m == 508u || m == 512u; +} + +class DispatchRouteRegistry { + public: + template + size_t register_group( + size_t dispatch_count, + const std::vector& ranges, + IsCompute&& is_compute) { + if (dispatch_count < owners_.size() || ranges.size() < 2) { + throw std::runtime_error("invalid WebGPU dispatch route group"); + } + + std::vector claimed(dispatch_count, false); + for (const auto& range : ranges) { + if (range.begin >= range.end || range.end > dispatch_count) { + throw std::runtime_error("invalid WebGPU dispatch route range"); + } + for (size_t i = range.begin; i < range.end; i++) { + if (!is_compute(i)) { + throw std::runtime_error( + "WebGPU dispatch route contains a copy command"); + } + if (claimed[i] || (i < owners_.size() && owners_[i] != kNoOwner)) { + throw std::runtime_error("overlapping WebGPU dispatch route ranges"); + } + claimed[i] = true; + } + } + + const size_t group = groups_.size(); + owners_.resize(dispatch_count, kNoOwner); + for (size_t i = 0; i < claimed.size(); i++) { + if (claimed[i]) { + owners_[i] = group; + } + } + groups_.push_back(ranges); + return group; + } + + template + void select( + size_t group, + size_t active_route, + const std::vector& active_grids, + SetGrid&& set_grid) const { + if (group >= groups_.size()) { + throw std::runtime_error("invalid WebGPU dispatch route group"); + } + const auto& ranges = groups_[group]; + if (active_route >= ranges.size()) { + throw std::runtime_error("invalid active WebGPU dispatch route"); + } + const auto& active = ranges[active_route]; + if (active_grids.size() != active.end - active.begin) { + throw std::runtime_error("WebGPU dispatch route grid count mismatch"); + } + for (const auto& grid : active_grids) { + if (grid.x == 0 || grid.y == 0) { + throw std::runtime_error( + "active WebGPU dispatch route has a zero grid"); + } + } + + for (const auto& range : ranges) { + for (size_t i = range.begin; i < range.end; i++) { + set_grid(i, {0, 0}); + } + } + for (size_t i = 0; i < active_grids.size(); i++) { + set_grid(active.begin + i, active_grids[i]); + } + } + + private: + static constexpr size_t kNoOwner = static_cast(-1); + std::vector> groups_; + std::vector owners_; +}; + // Given the workgroup count needed (1D) and the device's per-dimension // dispatch-count ceiling, compute a near-square 2D grid rather than // {max_dim, div_up(total, max_dim)} — maxing one dim pads the other with diff --git a/backends/webgpu/runtime/WebGPUExecutionOptions.cpp b/backends/webgpu/runtime/WebGPUExecutionOptions.cpp new file mode 100644 index 00000000000..19c0892030f --- /dev/null +++ b/backends/webgpu/runtime/WebGPUExecutionOptions.cpp @@ -0,0 +1,143 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include + +#include +#include +#include +#include + +namespace executorch::backends::webgpu { +namespace { + +thread_local WebGPUExecutionOptions execution_options; + +} // namespace + +WebGPUExecutionOptions current_webgpu_execution_options() { + return execution_options; +} + +ScopedWebGPUExecutionOptions::ScopedWebGPUExecutionOptions( + WebGPUExecutionOptions options) + : previous_(execution_options) { + execution_options = options; +} + +ScopedWebGPUExecutionOptions::~ScopedWebGPUExecutionOptions() { + execution_options = previous_; +} + +WebGPUExecutionPlan plan_webgpu_execution( + size_t dispatch_count, + size_t output_count, + ExecuteConfig config, + const std::vector& suppressible_outputs, + WebGPUGraphExecutionOptions options, + const std::vector& enabled_dispatches) { + if (!enabled_dispatches.empty() && + enabled_dispatches.size() != dispatch_count) { + throw std::runtime_error("WebGPU: enabled dispatch count mismatch"); + } + std::vector suppressed_dispatches(dispatch_count, false); + std::vector copy_outputs(output_count, true); + std::vector seen_output_ordinals(output_count, false); + + for (const auto& output : suppressible_outputs) { + if (output.output_ordinal >= output_count || + output.dispatch_begin >= output.dispatch_end || + output.dispatch_end > dispatch_count) { + throw std::runtime_error( + "WebGPU: invalid suppressible output range (output_id " + + std::to_string(output.output_id) + ")"); + } + if (seen_output_ordinals[output.output_ordinal]) { + throw std::runtime_error( + "WebGPU: duplicate suppressible output (output_id " + + std::to_string(output.output_id) + ")"); + } + seen_output_ordinals[output.output_ordinal] = true; + if (output.output_ordinal != options.suppress_output_ordinal) { + continue; + } + copy_outputs[output.output_ordinal] = false; + // Only the one ordinal matching suppress_output_ordinal reaches here (the + // duplicate check above rejects a repeat), so its dispatch range is + // disjoint by construction — mark it suppressed without a redundant overlap + // check. + for (size_t i = output.dispatch_begin; i < output.dispatch_end; i++) { + suppressed_dispatches[i] = true; + } + } + + WebGPUExecutionPlan plan; + plan.copy_outputs = std::move(copy_outputs); + + auto append_chunk = [&](size_t begin, size_t end) { + std::vector indices; + indices.reserve(end - begin); + for (size_t i = begin; i < end; i++) { + if (!suppressed_dispatches[i] && + (enabled_dispatches.empty() || enabled_dispatches[i])) { + indices.push_back(i); + } + } + if (!indices.empty()) { + plan.dispatch_chunks.push_back(std::move(indices)); + } + }; + + if (config.chunk_size == 0 || dispatch_count <= config.chunk_size) { + append_chunk(0, dispatch_count); + } else { + size_t start = 0; + size_t current_chunk = config.initial_chunk_size > 0 + ? config.initial_chunk_size + : config.chunk_size; + while (start < dispatch_count) { + const size_t end = std::min(start + current_chunk, dispatch_count); + append_chunk(start, end); + start = end; + current_chunk = config.chunk_size; + } + } + if (plan.dispatch_chunks.empty() && + std::any_of( + plan.copy_outputs.begin(), plan.copy_outputs.end(), [](bool copy) { + return copy; + })) { + plan.dispatch_chunks.emplace_back(); + } + return plan; +} + +WebGPUGraphExecutionOptions resolve_webgpu_graph_execution_options( + const std::vector& delegate_outputs, + WebGPUExecutionOptions options) { + if (options.discardable_output_data == nullptr) { + return {}; + } + if (!options.exact_method_certificate_verified) { + return {}; + } + + size_t match = kNoOutputOrdinal; + for (size_t i = 0; i < delegate_outputs.size(); i++) { + if (delegate_outputs[i] != options.discardable_output_data) { + continue; + } + if (match != kNoOutputOrdinal) { + return {}; + } + match = i; + } + return {match}; +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/WebGPUExecutionOptions.h b/backends/webgpu/runtime/WebGPUExecutionOptions.h new file mode 100644 index 00000000000..304f46839d8 --- /dev/null +++ b/backends/webgpu/runtime/WebGPUExecutionOptions.h @@ -0,0 +1,87 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include +#include +#include + +namespace executorch::backends::webgpu { + +constexpr size_t kNoOutputOrdinal = static_cast(-1); + +struct WebGPUExecutionOptions { + // The certificate must bind the exact PTE and method and prove one delegate, + // no portable nodes, and a unique leaf method output at this data pointer. + // The caller must keep this pointer valid and unchanged for the complete + // synchronous backend invocation in which these options are scoped. + const void* discardable_output_data = nullptr; + bool exact_method_certificate_verified = false; +}; + +struct WebGPUGraphExecutionOptions { + size_t suppress_output_ordinal = kNoOutputOrdinal; +}; + +struct ExecuteConfig { + size_t chunk_size = 0; + size_t initial_chunk_size = 0; +}; + +struct SuppressibleOutput { + int output_id = -1; + size_t output_ordinal = 0; + size_t dispatch_begin = 0; + size_t dispatch_end = 0; +}; + +struct WebGPUExecutionPlan { + std::vector> dispatch_chunks; + std::vector copy_outputs; +}; + +WebGPUExecutionPlan plan_webgpu_execution( + size_t dispatch_count, + size_t output_count, + ExecuteConfig config, + const std::vector& suppressible_outputs, + WebGPUGraphExecutionOptions options, + const std::vector& enabled_dispatches = {}); + +WebGPUGraphExecutionOptions resolve_webgpu_graph_execution_options( + const std::vector& delegate_outputs, + WebGPUExecutionOptions options); + +WebGPUExecutionOptions current_webgpu_execution_options(); + +class ScopedWebGPUExecutionOptions final { + public: + explicit ScopedWebGPUExecutionOptions(WebGPUExecutionOptions options); + ~ScopedWebGPUExecutionOptions(); + + ScopedWebGPUExecutionOptions(const ScopedWebGPUExecutionOptions&) = delete; + ScopedWebGPUExecutionOptions& operator=(const ScopedWebGPUExecutionOptions&) = + delete; + ScopedWebGPUExecutionOptions(ScopedWebGPUExecutionOptions&&) = delete; + ScopedWebGPUExecutionOptions& operator=(ScopedWebGPUExecutionOptions&&) = + delete; + + private: + WebGPUExecutionOptions previous_; +}; + +template +decltype(auto) with_webgpu_execution_options( + WebGPUExecutionOptions options, + Fn&& fn) { + ScopedWebGPUExecutionOptions scope(options); + return std::forward(fn)(); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/WebGPUGraph.cpp b/backends/webgpu/runtime/WebGPUGraph.cpp index e3f4b1e48d0..04331ee6ca7 100644 --- a/backends/webgpu/runtime/WebGPUGraph.cpp +++ b/backends/webgpu/runtime/WebGPUGraph.cpp @@ -7,20 +7,25 @@ */ #include +#include #include -#include -#include #include #include +#include #include #include +#include +#include +#include #include -#include +#include #include #include +#include +#include #include namespace executorch::backends::webgpu { @@ -30,9 +35,150 @@ namespace executorch::backends::webgpu { namespace { +const uint8_t* checked_inline_constant( + const uint8_t* data, + size_t data_size, + uint64_t offset, + size_t required_size, + const char* error_message) { + if (data == nullptr || offset > data_size || + required_size > data_size - static_cast(offset)) { + throw std::runtime_error(error_message); + } + return data + static_cast(offset); +} + +class ScopedBindGroupLayout final { + public: + explicit ScopedBindGroupLayout(WGPUBindGroupLayout handle) + : handle_(handle) {} + ~ScopedBindGroupLayout() { + if (handle_ != nullptr) { + wgpuBindGroupLayoutRelease(handle_); + } + } + ScopedBindGroupLayout(const ScopedBindGroupLayout&) = delete; + ScopedBindGroupLayout& operator=(const ScopedBindGroupLayout&) = delete; + + WGPUBindGroupLayout get() const { + return handle_; + } + + private: + WGPUBindGroupLayout handle_; +}; + +class ScopedBindGroup final { + public: + explicit ScopedBindGroup(WGPUBindGroup handle) : handle_(handle) {} + ~ScopedBindGroup() { + if (handle_ != nullptr) { + wgpuBindGroupRelease(handle_); + } + } + ScopedBindGroup(const ScopedBindGroup&) = delete; + ScopedBindGroup& operator=(const ScopedBindGroup&) = delete; + + WGPUBindGroup get() const { + return handle_; + } + + WGPUBindGroup release() { + WGPUBindGroup handle = handle_; + handle_ = nullptr; + return handle; + } + + private: + WGPUBindGroup handle_; +}; + +class ScopedComputePipeline final { + public: + explicit ScopedComputePipeline(WGPUComputePipeline handle) + : handle_(handle) {} + ~ScopedComputePipeline() { + if (handle_ != nullptr) { + wgpuComputePipelineRelease(handle_); + } + } + ScopedComputePipeline(const ScopedComputePipeline&) = delete; + ScopedComputePipeline& operator=(const ScopedComputePipeline&) = delete; + + WGPUComputePipeline release() { + WGPUComputePipeline handle = handle_; + handle_ = nullptr; + return handle; + } + + private: + WGPUComputePipeline handle_; +}; + +class ScopedComputePipelineRef final { + public: + explicit ScopedComputePipelineRef(WGPUComputePipeline handle) + : handle_(handle) { + wgpuComputePipelineAddRef(handle_); + } + ~ScopedComputePipelineRef() { + if (handle_ != nullptr) { + wgpuComputePipelineRelease(handle_); + } + } + ScopedComputePipelineRef(const ScopedComputePipelineRef&) = delete; + ScopedComputePipelineRef& operator=(const ScopedComputePipelineRef&) = delete; + + WGPUComputePipeline release() { + WGPUComputePipeline handle = handle_; + handle_ = nullptr; + return handle; + } + + private: + WGPUComputePipeline handle_; +}; + +void append_key_component(std::string& key, const std::string& value) { + const uint64_t size = value.size(); + key.append(reinterpret_cast(&size), sizeof(size)); + key.append(value); +} + +std::vector canonical_constants( + const WebGPUComputeDispatchDescriptor& descriptor) { + std::vector constants = descriptor.constants; + std::sort( + constants.begin(), + constants.end(), + [](const auto& left, const auto& right) { + return left.name < right.name; + }); + for (size_t i = 0; i < constants.size(); ++i) { + if (constants[i].name.empty()) { + throw std::runtime_error( + "WebGPU compute dispatch: empty specialization constant name"); + } + if (!std::isfinite(constants[i].value)) { + throw std::runtime_error( + "WebGPU compute dispatch: non-finite specialization constant"); + } + if (i > 0 && constants[i - 1].name == constants[i].name) { + throw std::runtime_error( + "WebGPU compute dispatch: duplicate specialization constant"); + } + if (constants[i].value == 0.0) { + constants[i].value = 0.0; + } + } + return constants; +} + // Op name the AOT exporter emits for a prepacked constant (must match the // serialized schema); compared in the prepack pre-scan below. constexpr const char* kPrepackOpName = "et_vk.prepack.default"; +constexpr const char* kQ4gswLinearOpName = "et_vk.linear_q4gsw.default"; +constexpr size_t kQ4gswOutputArg = 5; size_t vk_datatype_size(vkgraph::VkDataType dtype) { switch (dtype) { @@ -66,6 +212,31 @@ bool vk_datatype_is_int(vkgraph::VkDataType dtype) { } } +size_t storage_buffer_size(size_t nbytes) { + const size_t at_least_four = std::max(nbytes, size_t(4)); + if (at_least_four > std::numeric_limits::max() - 3u) { + throw std::runtime_error("WebGPU: storage buffer size overflows alignment"); + } + return (at_least_four + 3u) & ~size_t(3); +} + +void write_storage_buffer( + WGPUQueue queue, + WGPUBuffer buffer, + const void* data, + size_t nbytes) { + if (nbytes == 0u) { + return; + } + if (nbytes % 4u == 0u) { + wgpuQueueWriteBuffer(queue, buffer, 0, data, nbytes); + return; + } + std::vector padded(storage_buffer_size(nbytes), 0u); + std::memcpy(padded.data(), data, nbytes); + wgpuQueueWriteBuffer(queue, buffer, 0, padded.data(), padded.size()); +} + // Normalize a possibly-negative dim against rank; throws (fail-loud) if OOR. int normalize_dim(int dim, int rank, const char* op) { if (dim < 0) { @@ -78,27 +249,59 @@ int normalize_dim(int dim, int rank, const char* op) { return dim; } -// Uniform layout matching the fused-QKV WGSL Params struct (16B-aligned, 32B); -// identical to QuantizedLinear.cpp's Q4gswParams (kept local to this TU). -struct QkvFusedParams { - uint32_t M; - uint32_t N; - uint32_t K; - uint32_t K_packed; - uint32_t group_size; - uint32_t padded_N; - uint32_t has_bias; - uint32_t _pad; -}; -static_assert(sizeof(QkvFusedParams) == 32, "QkvFusedParams must be 32 bytes"); - } // namespace +std::string make_compute_pipeline_key( + const WebGPUComputeDispatchDescriptor& descriptor) { + if (descriptor.shader_name.empty()) { + throw std::runtime_error("WebGPU compute dispatch: empty shader name"); + } + if (descriptor.entry_point.empty()) { + throw std::runtime_error("WebGPU compute dispatch: empty entry point"); + } + + std::string key; + append_key_component(key, descriptor.shader_name); + append_key_component(key, descriptor.entry_point); + for (const auto& constant : canonical_constants(descriptor)) { + append_key_component(key, constant.name); + uint64_t value_bits = 0; + static_assert(sizeof(value_bits) == sizeof(constant.value)); + std::memcpy(&value_bits, &constant.value, sizeof(value_bits)); + key.append(reinterpret_cast(&value_bits), sizeof(value_bits)); + } + return key; +} + +void validate_compute_dispatch_descriptor( + const WebGPUComputeDispatchDescriptor& descriptor) { + (void)make_compute_pipeline_key(descriptor); + if (descriptor.bindings.empty()) { + throw std::runtime_error("WebGPU compute dispatch: no buffer bindings"); + } + for (const auto& binding : descriptor.bindings) { + if (binding.buffer == nullptr) { + throw std::runtime_error("WebGPU compute dispatch: null buffer binding"); + } + if (binding.size == 0) { + throw std::runtime_error("WebGPU compute dispatch: zero-size binding"); + } + if (binding.offset > UINT64_MAX - binding.size) { + throw std::runtime_error( + "WebGPU compute dispatch: binding range overflow"); + } + if (binding.offset + binding.size > wgpuBufferGetSize(binding.buffer)) { + throw std::runtime_error( + "WebGPU compute dispatch: binding range exceeds buffer"); + } + } +} + WebGPUGraph::WebGPUGraph() = default; WGPUBuffer WebGPUGraph::create_scratch_buffer(size_t nbytes) { WGPUBufferDescriptor buf_desc = {}; - buf_desc.size = nbytes > 0 ? nbytes : 4; + buf_desc.size = storage_buffer_size(nbytes); buf_desc.usage = WGPUBufferUsage_Storage | WGPUBufferUsage_CopyDst | WGPUBufferUsage_CopySrc; buf_desc.mappedAtCreation = false; @@ -108,7 +311,7 @@ WGPUBuffer WebGPUGraph::create_scratch_buffer(size_t nbytes) { } WGPUBuffer WebGPUGraph::acquire_scratch(size_t nbytes) { - nbytes = nbytes > 0 ? nbytes : 4; + nbytes = storage_buffer_size(nbytes); // Best-fit reuse: smallest free slot with size in [nbytes, 2*nbytes] -- the // 2x cap stops a large Cmax-sized buffer from backing a tiny request. Never // reuse an in_use slot (co-live safety). @@ -151,18 +354,156 @@ void WebGPUGraph::release_scratch(WGPUBuffer buffer) { } WGPUBuffer WebGPUGraph::make_uniform_buffer(const void* data, size_t size) { + if (data == nullptr || size == 0u) { + throw std::runtime_error("WebGPU: invalid uniform buffer data"); + } WGPUBufferDescriptor desc = {}; desc.size = size; desc.usage = WGPUBufferUsage_Uniform | WGPUBufferUsage_CopyDst; desc.mappedAtCreation = true; WGPUBuffer buffer = wgpuDeviceCreateBuffer(device_, &desc); + if (buffer == nullptr) { + throw std::runtime_error("WebGPU: failed to create uniform buffer"); + } void* mapped = wgpuBufferGetMappedRange(buffer, 0, size); + if (mapped == nullptr) { + wgpuBufferRelease(buffer); + throw std::runtime_error("WebGPU: failed to map uniform buffer"); + } std::memcpy(mapped, data, size); wgpuBufferUnmap(buffer); uniform_buffer_bytes_ += size; return buffer; } +size_t WebGPUGraph::add_compute_dispatch( + const WebGPUComputeDispatchDescriptor& descriptor) { + validate_compute_dispatch_descriptor(descriptor); + const WebGPUShaderInfo& shader_info = + get_webgpu_shader_info(descriptor.shader_name); + WGPUShaderModule shader = + get_or_create_shader(descriptor.shader_name, shader_info.source); + + const std::string pipeline_key = make_compute_pipeline_key(descriptor); + WGPUComputePipeline pipeline = nullptr; + auto pipeline_it = pipeline_cache_.find(pipeline_key); + if (pipeline_it != pipeline_cache_.end()) { + pipeline = pipeline_it->second; + } else { + const auto constants = canonical_constants(descriptor); + std::vector entries(constants.size()); + for (size_t i = 0; i < constants.size(); ++i) { + entries[i].key = {constants[i].name.data(), constants[i].name.size()}; + entries[i].value = constants[i].value; + } + + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.layout = nullptr; + pipeline_desc.compute.module = shader; + pipeline_desc.compute.entryPoint = { + descriptor.entry_point.data(), descriptor.entry_point.size()}; + pipeline_desc.compute.constantCount = entries.size(); + pipeline_desc.compute.constants = entries.data(); + ScopedComputePipeline created_pipeline( + wgpuDeviceCreateComputePipeline(device_, &pipeline_desc)); + pipeline = created_pipeline.release(); + if (pipeline == nullptr) { + throw std::runtime_error("WebGPU: failed to create compute pipeline"); + } + ScopedComputePipeline pipeline_owner(pipeline); + pipeline_cache_.emplace(pipeline_key, pipeline); + pipeline_owner.release(); + } + + ScopedBindGroupLayout layout( + wgpuComputePipelineGetBindGroupLayout(pipeline, 0)); + if (layout.get() == nullptr) { + throw std::runtime_error("WebGPU: failed to get bind-group layout"); + } + + std::vector entries(descriptor.bindings.size()); + for (size_t i = 0; i < descriptor.bindings.size(); ++i) { + entries[i].binding = i; + entries[i].buffer = descriptor.bindings[i].buffer; + entries[i].offset = descriptor.bindings[i].offset; + entries[i].size = descriptor.bindings[i].size; + } + WGPUBindGroupDescriptor bind_group_desc = {}; + bind_group_desc.layout = layout.get(); + bind_group_desc.entryCount = entries.size(); + bind_group_desc.entries = entries.data(); + ScopedBindGroup bind_group( + wgpuDeviceCreateBindGroup(device_, &bind_group_desc)); + if (bind_group.get() == nullptr) { + throw std::runtime_error("WebGPU: failed to create bind group"); + } + + ScopedComputePipelineRef dispatch_pipeline(pipeline); + const size_t dispatch_index = add_dispatch( + {pipeline, + bind_group.get(), + descriptor.grid.x, + descriptor.kernel_name.empty() ? descriptor.shader_name + : descriptor.kernel_name, + descriptor.grid.y}); + bind_group.release(); + dispatch_pipeline.release(); + return dispatch_index; +} + +size_t WebGPUGraph::add_dynamic_compute_dispatch_impl( + const WebGPUComputeDispatchDescriptor& descriptor, + int trigger_tensor_id, + std::function pick_grid) { + if (trigger_tensor_id < 0 || trigger_tensor_id >= num_values() || + get_value_type(trigger_tensor_id) != ValueType::Tensor) { + throw std::runtime_error( + "WebGPU dynamic dispatch: trigger must be a Tensor"); + } + if (!pick_grid) { + throw std::runtime_error("WebGPU dynamic dispatch: null grid picker"); + } + + const WebGPUDispatchGrid initial_grid = pick_grid(*this); + if (initial_grid.x == 0 || initial_grid.y == 0) { + throw std::runtime_error("WebGPU dynamic dispatch: zero grid"); + } + + WebGPUComputeDispatchDescriptor initial_descriptor = descriptor; + initial_descriptor.grid = initial_grid; + + // Reserve both vectors before creating GPU objects, then stage the sidecar. + // If GPU-object creation fails, removing the sidecar restores the graph; no + // operation that can fail remains after add_compute_dispatch succeeds. + const size_t new_size = dynamic_dispatch_grids_.size() + 1; + dynamic_dispatch_grids_.reserve(new_size); + pending_dynamic_dispatch_grids_.reserve(new_size); + + const size_t expected_index = dispatches_.size(); + dynamic_dispatch_grids_.push_back( + {expected_index, trigger_tensor_id, std::move(pick_grid)}); + try { + add_compute_dispatch(initial_descriptor); + } catch (...) { + dynamic_dispatch_grids_.pop_back(); + throw; + } + return expected_index; +} + +void WebGPUGraph::validate_dynamic_dispatch_route_ranges( + const std::vector& ranges) const { + for (const auto& dynamic_grid : dynamic_dispatch_grids_) { + for (const auto& range : ranges) { + if (range.begin <= dynamic_grid.dispatch_index && + dynamic_grid.dispatch_index < range.end) { + throw std::runtime_error( + "WebGPU dispatch cannot have both dynamic-grid and route ownership"); + } + } + } +} + void WebGPUGraph::update_symints_from_inputs( const std::vector& inputs) { for (const auto& src : symint_sources_) { @@ -207,7 +548,13 @@ void WebGPUGraph::update_symints_from_inputs( // elem_size (buffer-derived) would misread int64 host data as int32. int32_t val; if (inputs[pos].host_is_int64) { - val = static_cast(static_cast(host)[offset]); + const int64_t raw = static_cast(host)[offset]; + if (raw < std::numeric_limits::min() || + raw > std::numeric_limits::max()) { + throw std::runtime_error( + "select_as_symint: selected value is outside int32 range"); + } + val = static_cast(raw); } else { val = static_cast(host)[offset]; } @@ -293,11 +640,41 @@ void WebGPUGraph::propagate_resize() { pass++) { std::unordered_set processing; processing.swap(dirty_tensors_); - for (auto& hook : tensor_resize_hooks_) { - if (processing.count(hook.trigger_tensor_id) != 0) { - hook.fn(*this); + pending_dynamic_dispatch_grids_.clear(); + try { + for (auto& hook : tensor_resize_hooks_) { + if (processing.count(hook.trigger_tensor_id) != 0) { + hook.fn(*this); + } + } + + // A hook or picker may fail, so compute and validate every affected grid + // before changing any dispatch. The graph-owned staging vector has + // capacity for every registered dynamic grid and is reused on execute. + for (const auto& dynamic_grid : dynamic_dispatch_grids_) { + if (processing.count(dynamic_grid.trigger_tensor_id) == 0) { + continue; + } + const WebGPUDispatchGrid grid = dynamic_grid.pick_grid(*this); + if (grid.x == 0 || grid.y == 0) { + throw std::runtime_error("WebGPU dynamic dispatch: zero grid"); + } + pending_dynamic_dispatch_grids_.push_back( + {dynamic_grid.dispatch_index, grid}); } + } catch (...) { + pending_dynamic_dispatch_grids_.clear(); + // Keep both the current triggers and any cascaded outputs dirty so the + // caller can fix the hook or picker and retry without rebuilding. + dirty_tensors_.insert(processing.begin(), processing.end()); + throw; } + for (const auto& pending : pending_dynamic_dispatch_grids_) { + auto& dispatch = dispatches_[pending.dispatch_index]; + dispatch.workgroup_count_x = pending.grid.x; + dispatch.workgroup_count_y = pending.grid.y; + } + pending_dynamic_dispatch_grids_.clear(); } if (!dirty_tensors_.empty()) { throw std::runtime_error( @@ -375,9 +752,9 @@ WebGPUGraph::~WebGPUGraph() { void WebGPUGraph::build( const void* flatbuffer_data, const uint8_t* constant_data, + size_t constant_data_size, const executorch::runtime::NamedDataMap* named_data_map, - bool f16_kv_cache, - bool f16_accumulate_gemm) { + WebGPUGraphConfig config) { if (!device_) { auto* ctx = get_default_webgpu_context(); if (ctx) { @@ -396,16 +773,15 @@ void WebGPUGraph::build( // .pte byte sources for prepack-time constant materialization (build-only). constant_data_ = constant_data; + constant_data_size_ = constant_data_size; named_data_map_ = named_data_map; // f16 KV cache (runtime opt-in): store K/V caches as f16 iff the opt-in is // set AND the device negotiated shader-f16 (fail-closed). + config_ = config; const WebGPUContext* kv_ctx = get_default_webgpu_context(); - kv_f16_ = f16_kv_cache && (kv_ctx != nullptr && kv_ctx->shader_f16_supported); - - // f16-accumulate q4gsw steel prefill GEMM (runtime opt-in). QuantizedLinear - // additionally gates the kernel on the negotiated shader-f16 feature. - f16_accumulate_gemm_ = f16_accumulate_gemm; + kv_f16_ = config_.f16_kv_cache && + (kv_ctx != nullptr && kv_ctx->shader_f16_supported); // Phase 1: Create all values const auto* values = graph->values(); @@ -435,6 +811,12 @@ void WebGPUGraph::build( if (!a) { continue; } + if (oc->name()->str() == "sym_size.int" && a->size() >= 3 && values) { + const auto* out = values->Get(a->Get(2)); + if (out && out->value_type() == vkgraph::GraphTypes::SymInt) { + dynamic_tensor_ids_.insert(static_cast(a->Get(0))); + } + } // f16 KV: tag sdpa K/V cache values (args[3],[4]) for half-size alloc. // Inert unless kv_f16_ (runtime opt-in) is set. if (kv_f16_ && a->size() > 4 && @@ -476,11 +858,27 @@ void WebGPUGraph::build( continue; } for (unsigned j = 0; j < a->size(); j++) { - if (kv_cache_ids_.count(static_cast(a->Get(j))) != 0) { + const int id = static_cast(a->Get(j)); + if (kv_cache_ids_.count(id) != 0) { throw std::runtime_error( "WebGPU f16 KV: cache tensor consumed by non-sdpa op '" + nm + "' would misread the f16 buffer"); } + const auto* value = values ? values->Get(id) : nullptr; + if (value && value->value_type() == vkgraph::GraphTypes::ValueList) { + const auto* items = value->value_as_ValueList()->items(); + if (!items) { + continue; + } + for (unsigned k = 0; k < items->size(); k++) { + if (kv_cache_ids_.count(static_cast(items->Get(k))) != 0) { + throw std::runtime_error( + "WebGPU f16 KV: cache tensor consumed through a ValueList " + "by non-sdpa op '" + + nm + "' would misread the f16 buffer"); + } + } + } } } } @@ -502,12 +900,22 @@ void WebGPUGraph::build( size_t numel = 1; if (dims) { for (unsigned j = 0; j < dims->size(); j++) { - tensor.dims.push_back(static_cast(dims->Get(j))); - numel *= dims->Get(j); + const uint32_t dim = dims->Get(j); + tensor.dims.push_back(static_cast(dim)); + if (dim != 0 && numel > std::numeric_limits::max() / dim) { + throw std::runtime_error( + "WebGPU: tensor element count overflows"); + } + numel *= dim; } } tensor.elem_size = vk_datatype_size(vk_tensor->datatype()); + if (tensor.elem_size != 0 && + numel > std::numeric_limits::max() / tensor.elem_size) { + throw std::runtime_error("WebGPU: tensor byte size overflows"); + } tensor.is_int = vk_datatype_is_int(vk_tensor->datatype()); + tensor.is_bool = vk_tensor->datatype() == vkgraph::VkDataType::BOOL; tensor.is_int8 = vk_tensor->datatype() == vkgraph::VkDataType::INT8; tensor.nbytes = numel * tensor.elem_size; // Live dims start == max (serialized upper bound); resize_input shrinks @@ -519,16 +927,77 @@ void WebGPUGraph::build( // zero-initializes freshly-created buffers, so no explicit clear is // needed. Inert unless kv_f16_ (runtime opt-in) is set. if (kv_f16_ && kv_cache_ids_.count(i) != 0) { + if (tensor.is_int || tensor.elem_size != sizeof(float) || + tensor.nbytes != numel * sizeof(float)) { + throw std::runtime_error( + "WebGPU f16 KV: serialized cache tensor must be fp32"); + } tensor.elem_size = 2; tensor.nbytes = numel * 2; tensor.cur_nbytes = tensor.nbytes; tensor_mem_obj_ids_[i] = -1; WGPUBufferDescriptor buf_desc = {}; - buf_desc.size = std::max(tensor.nbytes, size_t(4)); + buf_desc.size = storage_buffer_size(tensor.nbytes); buf_desc.usage = WGPUBufferUsage_Storage | WGPUBufferUsage_CopyDst | WGPUBufferUsage_CopySrc; buf_desc.mappedAtCreation = false; tensor.buffer = wgpuDeviceCreateBuffer(device_, &buf_desc); + + // Mutable caches normally start empty. If the serialized graph owns + // an initialized cache constant, preserve it while changing storage + // representation instead of silently replacing it with zeros. + const int cache_constant_id = vk_tensor->constant_id(); + if (cache_constant_id >= 0) { + const auto* constants = graph->constants(); + if (!constants || + cache_constant_id >= static_cast(constants->size())) { + throw std::runtime_error( + "WebGPU f16 KV: cache constant id is out of range"); + } + const auto* bytes = constants->Get(cache_constant_id); + auto write_fp16_cache = [&](const uint8_t* src) { + std::vector converted(numel); + for (size_t e = 0; e < numel; e++) { + float value = 0.0f; + std::memcpy(&value, src + e * sizeof(float), sizeof(float)); + converted[e] = executorch::runtime::etensor::Half(value); + } + write_storage_buffer( + queue_, + tensor.buffer, + converted.data(), + converted.size() * sizeof(converted[0])); + }; + if (bytes->offset() != UINT64_MAX) { + write_fp16_cache(checked_inline_constant( + constant_data_, + constant_data_size_, + bytes->offset(), + numel * sizeof(float), + "WebGPU f16 KV: inline cache constant exceeds constant " + "data")); + } else if ( + bytes->named_key() != nullptr && named_data_map_ != nullptr) { + const std::string key = bytes->named_key()->str(); + auto data = named_data_map_->get_data(key.c_str()); + if (!data.ok()) { + throw std::runtime_error( + "WebGPU f16 KV: named cache constant '" + key + + "' not found"); + } + if (data->size() < numel * sizeof(float)) { + data->Free(); + throw std::runtime_error( + "WebGPU f16 KV: named cache constant '" + key + + "' is undersized"); + } + write_fp16_cache(static_cast(data->data())); + data->Free(); + } else { + throw std::runtime_error( + "WebGPU f16 KV: cache constant has no readable source"); + } + } break; } @@ -569,7 +1038,7 @@ void WebGPUGraph::build( prepack_src_ids.count(i) != 0 && direct_use_ids.count(i) == 0; if (!defer) { WGPUBufferDescriptor buf_desc = {}; - buf_desc.size = std::max(tensor.nbytes, size_t(4)); + buf_desc.size = storage_buffer_size(tensor.nbytes); buf_desc.usage = WGPUBufferUsage_Storage | WGPUBufferUsage_CopyDst | WGPUBufferUsage_CopySrc; buf_desc.mappedAtCreation = false; @@ -666,7 +1135,7 @@ void WebGPUGraph::build( shared_buffers_.resize(shared_buffer_sizes_.size(), nullptr); for (size_t id = 0; id < shared_buffer_sizes_.size(); id++) { WGPUBufferDescriptor buf_desc = {}; - buf_desc.size = std::max(shared_buffer_sizes_[id], size_t(4)); + buf_desc.size = storage_buffer_size(shared_buffer_sizes_[id]); buf_desc.usage = WGPUBufferUsage_Storage | WGPUBufferUsage_CopyDst | WGPUBufferUsage_CopySrc; buf_desc.mappedAtCreation = false; @@ -694,7 +1163,7 @@ void WebGPUGraph::build( // Create staging buffer for output readback WGPUBufferDescriptor staging_desc = {}; - staging_desc.size = std::max(tensors_[oid].nbytes, size_t(4)); + staging_desc.size = storage_buffer_size(tensors_[oid].nbytes); staging_desc.usage = WGPUBufferUsage_MapRead | WGPUBufferUsage_CopyDst; staging_desc.mappedAtCreation = false; output_staging_buffers_.push_back( @@ -710,349 +1179,49 @@ void WebGPUGraph::build( tensors_[oid].nbytes}); } - // Phase 3: Build operator dispatch chain - const auto* chain = graph->chain(); - - // QKV-concat fusion detection (auto-applied graph pass, no flag): the maps - // stay empty when no q/k/v triple matches -> the Phase-3 loop below runs - // verbatim. Find each attention q/k/v triple: EXACTLY 3 et_vk.linear_q4gsw - // ops sharing args[0] (the same input activation), in chain order q,k,v with - // the N-pattern {2048,512,512}, on the steel route (K%16==0), - // group_size%16==0, no bias. The fused kernel needs shader-f16 + a 256-thread - // WG, so gate on those (else leave the triple to the normal per-linear - // handlers). qkv_fused_skip holds all 3 op indices; qkv_anchor maps the FIRST - // op index -> its group, so the fused dispatch is emitted IN-PLACE at the - // anchor (correct execution order). - std::vector qkv_groups; - std::unordered_map - qkv_first; // first triple op -> group (repoint buffers) - std::unordered_map - qkv_last; // last triple op -> group (emit fused) - std::unordered_map - qkv_member; // any triple op -> group (record dispatch) - if (chain) { - bool device_ok = false; - { - WGPULimits limits = {}; - const bool have = - wgpuDeviceGetLimits(device_, &limits) == WGPUStatus_Success; - bool f16 = false; - if (auto* ctx = get_default_webgpu_context()) { - f16 = ctx->shader_f16_supported; - } - device_ok = - have && f16 && limits.maxComputeInvocationsPerWorkgroup >= 256u; - } - if (device_ok) { - // Group linear_q4gsw op indices by input id, preserving chain order. - std::unordered_map> by_input; - std::vector input_order; - for (unsigned i = 0; i < chain->size(); i++) { - const auto* oc = chain->Get(i); - if (oc->name()->str() != "et_vk.linear_q4gsw.default") { - continue; - } - const auto* a = oc->args(); - if (!a || a->size() < 6) { - continue; - } - const int inp = static_cast(a->Get(0)); - if (by_input.find(inp) == by_input.end()) { - input_order.push_back(inp); - } - by_input[inp].push_back(i); - } - auto op_arg = [&](unsigned oi, unsigned j) { - return static_cast(chain->Get(oi)->args()->Get(j)); - }; - for (int inp : input_order) { - const auto& ops = by_input[inp]; - if (ops.size() != 3) { - continue; // gate+up is a 2-group; o/down/lm_head are 1 each. - } - // args: [in, weight, scales, group_size, bias, out]. - const int wq = op_arg(ops[0], 1), sqid = op_arg(ops[0], 2), - bq = op_arg(ops[0], 4), oq = op_arg(ops[0], 5); - const int wk = op_arg(ops[1], 1), skid = op_arg(ops[1], 2), - bk = op_arg(ops[1], 4), ok = op_arg(ops[1], 5); - const int wv = op_arg(ops[2], 1), svid = op_arg(ops[2], 2), - bv = op_arg(ops[2], 4), ov = op_arg(ops[2], 5); - const int gsid = op_arg(ops[0], 3); - if (op_arg(ops[1], 3) != gsid || op_arg(ops[2], 3) != gsid) { - continue; // all 3 must share the group_size scalar. - } - if (get_value_type(bq) == ValueType::Tensor || - get_value_type(bk) == ValueType::Tensor || - get_value_type(bv) == ValueType::Tensor) { - continue; // fused kernel path assumes has_bias == 0. - } - const auto& twq = tensors_[wq]; - const auto& twk = tensors_[wk]; - const auto& twv = tensors_[wv]; - if (twq.dims.size() != 2 || twk.dims.size() != 2 || - twv.dims.size() != 2) { - continue; - } - const uint32_t Nq = static_cast(twq.dims[0]); - const uint32_t Nk = static_cast(twk.dims[0]); - const uint32_t Nv = static_cast(twv.dims[0]); - if (Nq != 2048u || Nk != 512u || Nv != 512u) { - continue; // kernel hardcodes N_Q=2048, N_KV=512 (Llama-3.2 GQA). - } - const uint32_t K_packed = static_cast(twq.dims[1]); - if (static_cast(twk.dims[1]) != K_packed || - static_cast(twv.dims[1]) != K_packed) { - continue; - } - const auto& tin = tensors_[inp]; - if (tin.dims.empty()) { - continue; - } - const uint32_t K = static_cast(tin.dims.back()); - if (K == 0 || K % 16u != 0u || K_packed != (K + 1u) / 2u) { - continue; // steel route stages a full BK=16 K-tile with no K-mask. - } - if (get_value_type(gsid) != ValueType::Int) { - continue; - } - const int64_t gsv = get_int(gsid); - if (gsv <= 0 || static_cast(gsv) % 16u != 0u) { - continue; // hoisted scale must be constant across the BK tile. - } - const uint32_t gs = static_cast(gsv); - const auto& tsq = tensors_[sqid]; - const auto& tsk = tensors_[skid]; - const auto& tsv = tensors_[svid]; - if (tsq.dims.size() != 2 || tsk.dims.size() != 2 || - tsv.dims.size() != 2) { - continue; - } - const uint32_t num_groups = static_cast(tsq.dims[0]); - if (static_cast(tsk.dims[0]) != num_groups || - static_cast(tsv.dims[0]) != num_groups) { - continue; - } - const uint32_t pNq = static_cast(tsq.dims[1]); - const uint32_t pNk = static_cast(tsk.dims[1]); - const uint32_t pNv = static_cast(tsv.dims[1]); - if (pNq < Nq || pNk < Nk || pNv < Nv || - num_groups < (K + gs - 1u) / gs) { - continue; - } - // All source + destination buffers must be live (Phase 1/2 allocated). - if (!twq.buffer || !twk.buffer || !twv.buffer || !tsq.buffer || - !tsk.buffer || !tsv.buffer || !tin.buffer || !tensors_[oq].buffer || - !tensors_[ok].buffer || !tensors_[ov].buffer) { - continue; - } + std::vector swiglu_fusions; + std::unordered_map swiglu_gate_producers; + std::unordered_map swiglu_anchors; + std::unordered_set swiglu_skipped_ops; + std::unordered_set claimed_fusion_ops; - QkvFusionGroup grp; - grp.input_id = inp; - grp.out_q = oq; - grp.out_k = ok; - grp.out_v = ov; - grp.weight_q = wq; - grp.weight_k = wk; - grp.weight_v = wv; - grp.scales_q = sqid; - grp.scales_k = skid; - grp.scales_v = svid; - grp.Nq = Nq; - grp.Nk = Nk; - grp.Nv = Nv; - grp.K = K; - grp.K_packed = K_packed; - grp.group_size = gs; - grp.num_groups = num_groups; - grp.padded_N_q = pNq; - grp.padded_N_k = pNk; - grp.padded_N_v = pNv; - grp.op_idx[0] = ops[0]; - grp.op_idx[1] = ops[1]; - grp.op_idx[2] = ops[2]; - const size_t gidx = qkv_groups.size(); - qkv_groups.push_back(grp); - qkv_first[ops[0]] = gidx; - qkv_last[ops[2]] = gidx; - qkv_member[ops[0]] = gidx; - qkv_member[ops[1]] = gidx; - qkv_member[ops[2]] = gidx; - } - } - } + std::vector qkv_fusions; + std::unordered_map qkv_first_ops; + std::unordered_map qkv_last_ops; + std::unordered_map qkv_member_ops; - // SwiGLU fusion detection (auto-applied graph pass, no flag): all sets stay - // empty when no SiLU-gate triple matches - // -> the Phase-3 loop below runs verbatim. Fold each - // SiLU-gate MLP triple sigmoid(g) -> mul(g,sig)=silu -> mul(silu,up)=out - // into ONE elementwise dispatch that computes sigmoid + silu in registers - // (gate + up read once, one output written): 8 traffic units -> 3. Bit-exact - // (same fp op order, and the sigmoid form matches sigmoid.wgsl). swiglu_skip - // holds the sigmoid + the 1st-mul op indices (their dispatches are dropped -- - // sig/silu become dead), and swiglu_anchor maps the 2nd-mul op (where out + - // up are both live) -> its group, so the fused dispatch is emitted IN-PLACE - // there (correct execution order). The sig/silu intermediates must be - // single-consumer (folding them can't strand a second reader). - std::vector> swiglu_groups; // {gate, up, out} - std::unordered_set swiglu_skip; // sigmoid + 1st-mul op indices - std::unordered_map swiglu_anchor; // 2nd-mul op idx -> group - // gate_proj op idx -> group: repoint gate to a PRIVATE pooled buffer there, - // before gate_proj is lowered (root-cause fix, see the detection guard - // below). - std::unordered_map swiglu_gate_acquire; - // out's last-use op idx -> group: release the pooled fused-output buffer - // after its final consumer's dispatch is built (pool recycling for the - // +memory). - std::unordered_map swiglu_out_release; - if (chain) { - // Consumer count: appearances as a NON-output (non-last) arg, ValueList- - // expanded. sig/silu are safe to fold only if each is consumed exactly once - // (by the 1st/2nd mul respectively) and nowhere else. - std::vector consumer_cnt(num_vals, 0); - for (unsigned i = 0; i < chain->size(); i++) { - const auto* a = chain->Get(i)->args(); - if (!a || a->size() == 0) { - continue; - } - for (unsigned j = 0; j + 1 < a->size(); j++) { - const int id = static_cast(a->Get(j)); - if (id < 0 || id >= num_vals) { - continue; - } - consumer_cnt[id]++; - if (value_types_[id] == ValueType::ValueList) { - for (int m : value_lists_[id]) { - if (m >= 0 && m < num_vals) { - consumer_cnt[m]++; - } - } - } - } - } - // Producer (op that writes each value = its last arg) + last-use (last op - // referencing a value in ANY arg, ValueList-expanded). Used to (a) find - // gate's producer op so gate can be repointed to a private buffer before it - // is written, and (b) find out's last consumer so the pooled out buffer is - // released only after it is truly dead. - std::vector producer(num_vals, -1); - std::vector last_use(num_vals, -1); - for (unsigned i = 0; i < chain->size(); i++) { - const auto* a = chain->Get(i)->args(); - if (!a || a->size() == 0) { - continue; - } - for (unsigned j = 0; j < a->size(); j++) { - const int id = static_cast(a->Get(j)); - if (id < 0 || id >= num_vals) { - continue; - } - last_use[id] = static_cast(i); - if (value_types_[id] == ValueType::ValueList) { - for (int m : value_lists_[id]) { - if (m >= 0 && m < num_vals) { - last_use[m] = static_cast(i); - } - } - } - } - const int outv = static_cast(a->Get(a->size() - 1)); - if (outv >= 0 && outv < num_vals) { - producer[outv] = static_cast(i); - } - } - struct SigInfo { - unsigned op; - int g_in; - }; - struct Mul1Info { - unsigned op; - int g_in; - unsigned sig_op; - int sig_out; - }; - std::unordered_map sigmoid_by_out; // sig_out -> {op, g} - std::unordered_map - mul1_by_out; // silu_out -> {op, g, sig...} - for (unsigned i = 0; i < chain->size(); i++) { - const auto* oc = chain->Get(i); - const std::string nm = oc->name()->str(); - const auto* a = oc->args(); - if (!a) { - continue; - } - if (nm == "aten.sigmoid.default" && a->size() >= 2) { - sigmoid_by_out[static_cast(a->Get(1))] = { - i, static_cast(a->Get(0))}; - continue; - } - if (nm != "aten.mul.Tensor" || a->size() < 3) { - continue; - } - const int x = static_cast(a->Get(0)); - const int y = static_cast(a->Get(1)); - const int out = static_cast(a->Get(2)); - // 2nd mul? one operand is a recorded silu (a 1st-mul output). - int silu = -1, up = -1; - if (mul1_by_out.count(x) != 0) { - silu = x; - up = y; - } else if (mul1_by_out.count(y) != 0) { - silu = y; - up = x; - } - if (silu >= 0) { - const Mul1Info& m1 = mul1_by_out[silu]; - const int g = m1.g_in; - const bool tensors_ok = g >= 0 && up >= 0 && out >= 0 && - get_value_type(g) == ValueType::Tensor && - get_value_type(up) == ValueType::Tensor && - get_value_type(out) == ValueType::Tensor; - if (tensors_ok) { - const auto& tg = tensors_[g]; - const auto& tu = tensors_[up]; - const auto& to = tensors_[out]; - // Elementwise, all fp32, identical dims (so the fused output's live - // dims - // == gate's on resize), live buffers, and single-consumer - // intermediates. gate must be consumed by EXACTLY the sigmoid + - // 1st-mul (consumer_cnt - // == 2) so it is dead once the fused dispatch reads it, and have a - // real producer op (so it can be repointed before it is written). - if (tg.buffer && tu.buffer && to.buffer && tg.elem_size == 4 && - tu.elem_size == 4 && to.elem_size == 4 && tg.dims == tu.dims && - tg.dims == to.dims && tg.nbytes == tu.nbytes && - tg.nbytes == to.nbytes && consumer_cnt[m1.sig_out] == 1 && - consumer_cnt[silu] == 1 && consumer_cnt[g] == 2 && - producer[g] >= 0) { - const size_t gidx = swiglu_groups.size(); - swiglu_groups.push_back({g, up, out}); - swiglu_skip.insert(m1.sig_op); - swiglu_skip.insert(m1.op); - swiglu_anchor[i] = gidx; - // The serialized memory planner reuse-aliases up onto gate's slot - // (gate dies at the 1st mul, up_proj is emitted between the muls), - // so up_proj would stomp gate's buffer before the fused dispatch - // (at this 2nd-mul anchor) reads it. Give gate a private buffer at - // its producer op; release it right after the fused read. out is - // likewise pooled + released after last_use[out] (its final - // consumer, e.g. down_proj). - swiglu_gate_acquire[static_cast(producer[g])] = gidx; - swiglu_out_release[static_cast(last_use[out])] = gidx; - } - } - continue; // a 2nd-mul is never also a 1st-mul - } - // 1st mul? inputs are exactly {sigmoid input g, sigmoid output sig}. - auto sx = sigmoid_by_out.find(x); - auto sy = sigmoid_by_out.find(y); - if (sx != sigmoid_by_out.end() && sx->second.g_in == y) { - mul1_by_out[out] = {i, y, sx->second.op, x}; - } else if (sy != sigmoid_by_out.end() && sy->second.g_in == x) { - mul1_by_out[out] = {i, x, sy->second.op, y}; - } - } - } + const auto* chain = graph->chain(); + passes::detect_qkv_bk64_fusions( + *this, + graph, + num_vals, + qkv_fusions, + qkv_first_ops, + qkv_last_ops, + qkv_member_ops); + passes::detect_swiglu_fusions( + *this, + graph, + num_vals, + swiglu_fusions, + swiglu_gate_producers, + swiglu_anchors, + swiglu_skipped_ops, + claimed_fusion_ops); + + // SwiGLU keeps precedence when the exact QKV geometry is also formed by a + // q projection plus gate/up projections. QKV detection runs first because it + // validates constant geometry, but it has no side effects until Phase 3; now + // discard candidates claimed by the completed SwiGLU pass and rebuild the + // index maps for the retained groups. + passes::retain_unclaimed_qkv_fusions( + qkv_fusions, + qkv_first_ops, + qkv_last_ops, + qkv_member_ops, + claimed_fusion_ops); + // Phase 3: Build operator dispatch chain if (chain) { for (unsigned i = 0; i < chain->size(); i++) { const auto* op_call = chain->Get(i); @@ -1070,89 +1239,75 @@ void WebGPUGraph::build( } } - // SwiGLU fusion. At gate_proj repoint gate to a private pooled buffer (so - // up_proj can't stomp its planner-aliased slot before the fused reads - // it); drop the folded sigmoid + 1st-mul; at the 2nd-mul anchor emit ONE - // fused silu*up dispatch then release gate. Sets empty when - // no SwiGLU triple matched (verbatim path). - { - auto ga = swiglu_gate_acquire.find(i); - if (ga != swiglu_gate_acquire.end()) { - // Repoint BEFORE gate_proj is lowered below, so it writes the private - // buffer. gate_proj falls through to normal lowering (not - // skip/anchor). - const int gate_id = swiglu_groups[ga->second][0]; - tensors_[gate_id].buffer = acquire_scratch(tensors_[gate_id].nbytes); - } - if (swiglu_skip.count(i) != 0) { - continue; // sigmoid / 1st-mul: folded into the fused dispatch - } - auto sa = swiglu_anchor.find(i); - if (sa != swiglu_anchor.end()) { - const auto& grp = swiglu_groups[sa->second]; - add_swiglu_fused_dispatch(grp[0], grp[1], grp[2]); - // gate is dead once the fused dispatch has read it - // (consumer_cnt[g]==2, both folded) -> return its buffer to the pool - // for the next layer. - release_scratch(tensors_[grp[0]].buffer); - continue; - } + const auto gate_it = swiglu_gate_producers.find(i); + if (gate_it != swiglu_gate_producers.end()) { + const int gate_id = swiglu_fusions[gate_it->second].gate_id; + tensors_[gate_id].buffer = acquire_scratch(tensors_[gate_id].nbytes); + } + const auto anchor_it = swiglu_anchors.find(i); + if (anchor_it != swiglu_anchors.end()) { + const passes::SwiGluFusion& fusion = swiglu_fusions[anchor_it->second]; + passes::add_silu_mul_fused_dispatch( + *this, + fusion.common_input_id, + fusion.gate_id, + fusion.up_id, + fusion.out_id); + release_scratch(tensors_[fusion.gate_id].buffer); + continue; + } + if (swiglu_skipped_ops.count(i) != 0) { + continue; } - // QKV fusion (M-gated): keep the 3 separate q/k/v linears AND add a fused - // multi-output GEMM; the fused resize hook selects by LIVE M (prefill M>1 - // -> fused runs, the 3 zeroed; decode M==1 -> the 3 coop4 GEMVs run, - // fused zeroed -- the fused 64x64 tile is ~4x slower than coop4 at M=1). - // At the FIRST triple op, repoint the 3 outputs to FRESH distinct - // buffers: the planner reuse-aliases q/k/v (each dies right after RoPE), - // which is fatal for a simultaneous fused write, so BOTH paths use - // non-aliased storage. All maps empty when no triple matches (verbatim - // path). - { - auto fit = qkv_first.find(i); - if (fit != qkv_first.end()) { - const auto& g = qkv_groups[fit->second]; - tensors_[g.out_q].buffer = - create_scratch_buffer(tensors_[g.out_q].nbytes); - tensors_[g.out_k].buffer = - create_scratch_buffer(tensors_[g.out_k].nbytes); - tensors_[g.out_v].buffer = - create_scratch_buffer(tensors_[g.out_v].nbytes); + + const auto qkv_first = qkv_first_ops.find(i); + if (qkv_first != qkv_first_ops.end()) { + passes::QkvBk64Fusion& fusion = qkv_fusions[qkv_first->second]; + for (int output_id : fusion.output_ids) { + tensors_[output_id].buffer = + create_scratch_buffer(tensors_[output_id].nbytes); } } + const size_t dispatch_begin = dispatches_.size(); webgpu_operator_registry().get_op_fn(op_name)(*this, args); - - { - auto mit = qkv_member.find(i); - if (mit != qkv_member.end()) { - QkvFusionGroup& g = qkv_groups[mit->second]; - const size_t di = num_dispatches() - 1; // this linear's dispatch - if (i == g.op_idx[0]) { - g.sep_dispatch[0] = di; - // Emit the fused dispatch RIGHT AFTER the q-linear (the anchor) so - // at M>1 it writes q/k/v BEFORE any consumer. q/k/v may be - // interleaved with rope in the chain, so emitting it at the LAST - // triple op would let a consumer (rope-q) read still-unwritten - // fresh_q -> garbage. - add_qkv_fused_dispatch(g); - } else if (i == g.op_idx[1]) { - g.sep_dispatch[1] = di; - } else { - g.sep_dispatch[2] = di; - } + const size_t dispatch_end = dispatches_.size(); + + const auto qkv_member = qkv_member_ops.find(i); + if (qkv_member != qkv_member_ops.end()) { + passes::QkvBk64Fusion& fusion = qkv_fusions[qkv_member->second]; + size_t member = 0; + while (member < 3 && fusion.op_indices[member] != i) { + member++; } - auto lit = qkv_last.find(i); - if (lit != qkv_last.end()) { - // All 3 sep dispatch indices + the fused index are now known. - add_qkv_fused_hook(qkv_groups[lit->second]); + if (member == 3 || dispatch_end <= dispatch_begin) { + throw std::runtime_error( + "linear_q4gsw_bk64_qkv: malformed member dispatch range"); } + fusion.separate_begin[member] = dispatch_begin; + fusion.separate_end[member] = dispatch_end; + if (member == 0) { + passes::add_qkv_bk64_dispatch(*this, fusion); + } + } + const auto qkv_last = qkv_last_ops.find(i); + if (qkv_last != qkv_last_ops.end()) { + passes::add_qkv_bk64_resize_hook(*this, qkv_fusions[qkv_last->second]); } - // SwiGLU: this op is out's last consumer (its dispatch just captured - // out's buffer above) -> return the pooled fused-output buffer for reuse. - { - auto orl = swiglu_out_release.find(i); - if (orl != swiglu_out_release.end()) { - release_scratch(tensors_[swiglu_groups[orl->second][2]].buffer); + + if (i + 1 == chain->size() && op_name == kQ4gswLinearOpName && + args.size() > kQ4gswOutputArg && dispatch_end > dispatch_begin) { + const int output_id = args[kQ4gswOutputArg]; + const auto output_it = + std::find(output_ids_.begin(), output_ids_.end(), output_id); + if (output_it != output_ids_.end() && + std::count(output_ids_.begin(), output_ids_.end(), output_id) == + 1) { + suppressible_outputs_.push_back( + {output_id, + static_cast(output_it - output_ids_.begin()), + dispatch_begin, + dispatch_end}); } } } @@ -1163,6 +1318,7 @@ void WebGPUGraph::build( // The .pte bytes are freed right after build() returns (WebGPUBackend // processed->Free()), so clear the build-only source pointers. constant_data_ = nullptr; + constant_data_size_ = 0; named_data_map_ = nullptr; } @@ -1174,15 +1330,18 @@ void WebGPUGraph::materialize_constant(int const_value_id, WGPUBuffer dst) { std::to_string(const_value_id)); } const ConstantSource& cs = it->second; - if (cs.nbytes == 0) { - return; - } if (cs.inline_offset != UINT64_MAX) { - if (constant_data_ == nullptr) { - throw std::runtime_error("WebGPU: inline constant data is null"); + const uint8_t* data = checked_inline_constant( + constant_data_, + constant_data_size_, + cs.inline_offset, + cs.nbytes, + "WebGPU: inline constant exceeds constant data"); + if (cs.nbytes != 0) { + write_storage_buffer(queue_, dst, data, cs.nbytes); } - wgpuQueueWriteBuffer( - queue_, dst, 0, constant_data_ + cs.inline_offset, cs.nbytes); + } else if (cs.nbytes == 0) { + return; } else if (!cs.named_key.empty() && named_data_map_ != nullptr) { auto buf = named_data_map_->get_data(cs.named_key.c_str()); if (!buf.ok()) { @@ -1193,7 +1352,7 @@ void WebGPUGraph::materialize_constant(int const_value_id, WGPUBuffer dst) { throw std::runtime_error( "WebGPU: named constant '" + cs.named_key + "' undersized"); } - wgpuQueueWriteBuffer(queue_, dst, 0, buf->data(), cs.nbytes); + write_storage_buffer(queue_, dst, buf->data(), cs.nbytes); buf->Free(); } else { throw std::runtime_error("WebGPU: constant has no source"); @@ -1215,8 +1374,16 @@ WGPUShaderModule WebGPUGraph::get_or_create_shader( WGPUShaderModuleDescriptor shader_desc = {}; shader_desc.nextInChain = &wgsl_desc.chain; WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device_, &shader_desc); + if (shader == nullptr) { + throw std::runtime_error("WebGPU: failed to create shader module"); + } - shader_cache_[key] = shader; + try { + shader_cache_.emplace(key, shader); + } catch (...) { + wgpuShaderModuleRelease(shader); + throw; + } return shader; } @@ -1258,408 +1425,6 @@ WGPUBindGroupLayout WebGPUGraph::get_or_create_bgl( return bgl; } -void WebGPUGraph::add_qkv_fused_dispatch(QkvFusionGroup& g) { - const uint32_t N = g.Nq + g.Nk + g.Nv; // fused output width (3072) - - const auto& in = tensors_[g.input_id]; - const auto& out_q = tensors_[g.out_q]; - const auto& out_k = tensors_[g.out_k]; - const auto& out_v = tensors_[g.out_v]; - const auto& wq = tensors_[g.weight_q]; - const auto& wk = tensors_[g.weight_k]; - const auto& wv = tensors_[g.weight_v]; - const auto& sq = tensors_[g.scales_q]; - const auto& sk = tensors_[g.scales_k]; - const auto& sv = tensors_[g.scales_v]; - - // Buffers were repointed to FRESH distinct slots at the first triple op (see - // the build() op-walk), so out_q/k/v no longer alias. Live M from the shared - // input. - uint64_t in_numel = 1; - for (int64_t d : in.dims) { - in_numel *= static_cast(d); - } - const uint32_t M = static_cast(in_numel / g.K); - - // Fused weight [N, K_packed]: a byte-contiguous row-stack of Wq;Wk;Wv (q4gsw - // packs each output row independently along a shared K_packed, so stacking - // along N is a flat append -- bit-exact). Fused scales [num_groups, N]: a - // strided PER-GROUP-ROW gather (dest row stride N != the per-linear source - // strides padded_N_{q,k,v}), NOT a flat append. Both dtor-freed via scratch. - const uint64_t kp = static_cast(g.K_packed); // packed bytes / row - const uint64_t fs = sizeof(float); - WGPUBuffer fused_weight = create_scratch_buffer(static_cast(N) * kp); - WGPUBuffer fused_scales = create_scratch_buffer( - static_cast(g.num_groups) * N * sizeof(float)); - - // Sources are direct constants materialized in Phase 1 (or prepack outputs - // materialized earlier in Phase 3); all writes are already enqueued on - // queue_, so this build-time copy sees the materialized bytes. - WGPUCommandEncoder enc = wgpuDeviceCreateCommandEncoder(device_, nullptr); - wgpuCommandEncoderCopyBufferToBuffer( - enc, wq.buffer, 0, fused_weight, 0, static_cast(g.Nq) * kp); - wgpuCommandEncoderCopyBufferToBuffer( - enc, - wk.buffer, - 0, - fused_weight, - static_cast(g.Nq) * kp, - static_cast(g.Nk) * kp); - wgpuCommandEncoderCopyBufferToBuffer( - enc, - wv.buffer, - 0, - fused_weight, - static_cast(g.Nq + g.Nk) * kp, - static_cast(g.Nv) * kp); - for (uint32_t grp = 0; grp < g.num_groups; grp++) { - const uint64_t dst_row = static_cast(grp) * N * fs; - wgpuCommandEncoderCopyBufferToBuffer( - enc, - sq.buffer, - static_cast(grp) * g.padded_N_q * fs, - fused_scales, - dst_row, - static_cast(g.Nq) * fs); - wgpuCommandEncoderCopyBufferToBuffer( - enc, - sk.buffer, - static_cast(grp) * g.padded_N_k * fs, - fused_scales, - dst_row + static_cast(g.Nq) * fs, - static_cast(g.Nk) * fs); - wgpuCommandEncoderCopyBufferToBuffer( - enc, - sv.buffer, - static_cast(grp) * g.padded_N_v * fs, - fused_scales, - dst_row + static_cast(g.Nq + g.Nk) * fs, - static_cast(g.Nv) * fs); - } - WGPUCommandBuffer cmd = wgpuCommandEncoderFinish(enc, nullptr); - wgpuQueueSubmit(queue_, 1, &cmd); - wgpuCommandBufferRelease(cmd); - wgpuCommandEncoderRelease(enc); - - // Params UBO (owned; rewritten by the resize hook). padded_N == N (fused - // scales row stride); has_bias == 0 (attention q/k/v are bias-less). - QkvFusedParams params = {}; - params.M = M; - params.N = N; - params.K = g.K; - params.K_packed = g.K_packed; - params.group_size = g.group_size; - params.padded_N = N; - params.has_bias = 0; - WGPUBufferDescriptor u_desc = {}; - u_desc.size = sizeof(QkvFusedParams); - u_desc.usage = WGPUBufferUsage_Uniform | WGPUBufferUsage_CopyDst; - u_desc.mappedAtCreation = true; - WGPUBuffer uniform_buffer = wgpuDeviceCreateBuffer(device_, &u_desc); - std::memcpy( - wgpuBufferGetMappedRange(uniform_buffer, 0, sizeof(QkvFusedParams)), - ¶ms, - sizeof(QkvFusedParams)); - wgpuBufferUnmap(uniform_buffer); - add_uniform_buffer_bytes(sizeof(QkvFusedParams)); - - // 4-byte dummy for the fixed bias binding (has_bias == 0). - WGPUBuffer bias_dummy = create_scratch_buffer(4); - - // Bespoke 8-binding layout: 3 rw-storage outputs + 4 ro-storage + 1 uniform. - // One-off shader/bgl/pipeline owned by the dispatch (matches - // q4gsw_linear_impl). - WGPUBindGroupLayoutEntry entries[8] = {}; - for (uint32_t i = 0; i < 3; i++) { - entries[i].binding = i; - entries[i].visibility = WGPUShaderStage_Compute; - entries[i].buffer.type = WGPUBufferBindingType_Storage; - } - for (uint32_t i = 3; i < 7; i++) { - entries[i].binding = i; - entries[i].visibility = WGPUShaderStage_Compute; - entries[i].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - } - entries[7].binding = 7; - entries[7].visibility = WGPUShaderStage_Compute; - entries[7].buffer.type = WGPUBufferBindingType_Uniform; - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 8; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device_, &bgl_desc); - - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kQ4gswLinearGemmQkvFusedWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device_, &shader_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device_, &pl_desc); - - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device_, &pipeline_desc); - - WGPUBindGroupEntry bg[8] = {}; - bg[0].binding = 0; - bg[0].buffer = out_q.buffer; - bg[0].size = out_q.nbytes; - bg[1].binding = 1; - bg[1].buffer = out_k.buffer; - bg[1].size = out_k.nbytes; - bg[2].binding = 2; - bg[2].buffer = out_v.buffer; - bg[2].size = out_v.nbytes; - bg[3].binding = 3; - bg[3].buffer = in.buffer; - bg[3].size = in.nbytes; - bg[4].binding = 4; - bg[4].buffer = fused_weight; - bg[4].size = static_cast(N) * kp; - bg[5].binding = 5; - bg[5].buffer = fused_scales; - bg[5].size = static_cast(g.num_groups) * N * fs; - bg[6].binding = 6; - bg[6].buffer = bias_dummy; - bg[6].size = 4; - bg[7].binding = 7; - bg[7].buffer = uniform_buffer; - bg[7].size = sizeof(QkvFusedParams); - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 8; - bg_desc.entries = bg; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device_, &bg_desc); - - // 1D dispatch over ceil(M/BM) * ceil(N/BN) tiles (BM=BN=64), matching the - // kernel's nbN = ceil(N/64) tile decode (NOT grid-strided). - const uint32_t nbN = (N + 63u) / 64u; - const uint32_t nbM = (M + 63u) / 64u; - const size_t fused_idx = - add_dispatch({pipeline, bind_group, nbN * nbM, "linear_q4gsw_qkv_fused"}); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); - own_uniform_buffer(uniform_buffer); - g.fused_dispatch = fused_idx; // consumed by add_qkv_fused_hook at the last op - g.fused_params = uniform_buffer; -} - -namespace { -// Uniform layout matching silu_mul_fused.wgsl Params (16B-aligned). -struct SiluMulParams { - uint32_t num_elements; - uint32_t _pad[3]; -}; -} // namespace - -// SwiGLU fusion: emit ONE elementwise dispatch computing -// out = (gate * sigmoid(gate)) * up, replacing the sigmoid + 2 muls. -// Elementwise (no M-gate: identical at decode and prefill). -void WebGPUGraph::add_swiglu_fused_dispatch( - int gate_id, - int up_id, - int out_id) { - // Private distinct output buffer (mirrors the QKV aliasing guard): the - // planner reuse-aliases `out` onto a dead slot (e.g. sigmoid's), which - // without this would bind the same buffer as ro `gate` AND rw `output` -> - // Dawn writable-aliasing / all-zeros. Repoint BEFORE the bind group so it - // captures the private buffer; downstream consumers (lowered later) also see - // it. gate is still in_use here (released only after this call), so - // acquire_scratch hands out a DISTINCT slot. Pooled (not dedicated): the - // caller releases it after out's last consumer, so N layers recycle a small - // constant of buffers. tensor_mem_obj_ids_[out] stays - // >= 0, so the dtor never per-tensor-frees it (scratch_pool_ owns it). - tensors_[out_id].buffer = acquire_scratch(tensors_[out_id].nbytes); - - const auto& gate = tensors_[gate_id]; - const auto& up = tensors_[up_id]; - const auto& out = tensors_[out_id]; - const uint32_t num_elements = - static_cast(out.nbytes / sizeof(float)); - - SiluMulParams params = {num_elements, {0u, 0u, 0u}}; - WGPUBufferDescriptor u_desc = {}; - u_desc.size = sizeof(SiluMulParams); - u_desc.usage = WGPUBufferUsage_Uniform | WGPUBufferUsage_CopyDst; - u_desc.mappedAtCreation = true; - WGPUBuffer uniform_buffer = wgpuDeviceCreateBuffer(device_, &u_desc); - std::memcpy( - wgpuBufferGetMappedRange(uniform_buffer, 0, sizeof(SiluMulParams)), - ¶ms, - sizeof(SiluMulParams)); - wgpuBufferUnmap(uniform_buffer); - add_uniform_buffer_bytes(sizeof(SiluMulParams)); - - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kSiluMulFusedWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device_, &shader_desc); - - // BGL: gate (ro) 0, up (ro) 1, output (rw) 2, params (uniform) 3. - WGPUBindGroupLayoutEntry entries[4] = {}; - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[1].binding = 1; - entries[1].visibility = WGPUShaderStage_Compute; - entries[1].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[2].binding = 2; - entries[2].visibility = WGPUShaderStage_Compute; - entries[2].buffer.type = WGPUBufferBindingType_Storage; - entries[3].binding = 3; - entries[3].visibility = WGPUShaderStage_Compute; - entries[3].buffer.type = WGPUBufferBindingType_Uniform; - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 4; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device_, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device_, &pl_desc); - - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device_, &pipeline_desc); - - WGPUBindGroupEntry bg[4] = {}; - bg[0].binding = 0; - bg[0].buffer = gate.buffer; - bg[0].size = gate.nbytes; - bg[1].binding = 1; - bg[1].buffer = up.buffer; - bg[1].size = up.nbytes; - bg[2].binding = 2; - bg[2].buffer = out.buffer; - bg[2].size = out.nbytes; - bg[3].binding = 3; - bg[3].buffer = uniform_buffer; - bg[3].size = sizeof(SiluMulParams); - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 4; - bg_desc.entries = bg; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device_, &bg_desc); - - const uint32_t wg = kSiluMulFusedWorkgroupSizeX; - const uint32_t workgroup_count = (num_elements + wg - 1) / wg; - if (workgroup_count > 65535) { - throw std::runtime_error( - "silu_mul_fused: workgroup count exceeds 65535 (1D dispatch limit)"); - } - add_dispatch({pipeline, bind_group, workgroup_count, "silu_mul_fused"}); - const size_t dispatch_idx = num_dispatches() - 1; - - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); - own_uniform_buffer(uniform_buffer); - - // Dynamic shapes: gate/up/out share dims, so out's live dims == gate's; - // recompute num_elements + dispatch from gate's live shape. Triggers on gate - // -- exactly the input the folded sigmoid's hook keyed on, so it is dirtied - // on every resize. - WGPUBuffer params_buf = uniform_buffer; - add_tensor_resize_hook( - gate_id, [gate_id, out_id, wg, dispatch_idx, params_buf](WebGPUGraph& g) { - const auto& d = g.cur_dims(gate_id); - g.set_cur_dims(out_id, d); - uint64_t numel = 1; - for (int64_t v : d) { - numel *= static_cast(v); - } - SiluMulParams p = {static_cast(numel), {0u, 0u, 0u}}; - wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); - g.dispatch_at(dispatch_idx).workgroup_count_x = - (static_cast(numel) + wg - 1) / wg; - }); -} - -// M-gate coordinator: registered at the LAST triple op (all dispatch indices -// known). Prefill (M>1): run the fused GEMM, zero the 3 separate linears. -// Decode (M==1): zero the fused, leave the 3 coop4 GEMVs (their own hooks set -// the decode wg) -- the fused 64x64 tile wastes 63/64 rows at M=1. Recomputes -// live M + the 3 output cur_dims + fused params. Inert on a static graph; a -// workgroup_count of 0 = no-op. -void WebGPUGraph::add_qkv_fused_hook(const QkvFusionGroup& g) { - const int input_id = g.input_id, out_q_id = g.out_q, out_k_id = g.out_k, - out_v_id = g.out_v; - const uint32_t K = g.K, Kp = g.K_packed, gs = g.group_size, Nq = g.Nq, - Nk = g.Nk, Nv = g.Nv, Nf = g.Nq + g.Nk + g.Nv; - const size_t fused_idx = g.fused_dispatch, sep0 = g.sep_dispatch[0], - sep1 = g.sep_dispatch[1], sep2 = g.sep_dispatch[2]; - WGPUBuffer params_buf = g.fused_params; - add_tensor_resize_hook( - input_id, - [input_id, - out_q_id, - out_k_id, - out_v_id, - K, - Kp, - gs, - Nq, - Nk, - Nv, - Nf, - fused_idx, - sep0, - sep1, - sep2, - params_buf](WebGPUGraph& gr) { - const auto& d = gr.cur_dims(input_id); - uint64_t numel = 1; - for (int64_t v : d) { - numel *= static_cast(v); - } - const uint32_t m = static_cast(numel / K); - std::vector oq = d; - oq.back() = static_cast(Nq); - std::vector ok = d; - ok.back() = static_cast(Nk); - std::vector ov = d; - ov.back() = static_cast(Nv); - gr.set_cur_dims(out_q_id, oq); - gr.set_cur_dims(out_k_id, ok); - gr.set_cur_dims(out_v_id, ov); - QkvFusedParams p = {}; - p.M = m; - p.N = Nf; - p.K = K; - p.K_packed = Kp; - p.group_size = gs; - p.padded_N = Nf; - p.has_bias = 0; - wgpuQueueWriteBuffer(gr.queue(), params_buf, 0, &p, sizeof(p)); - if (m > 1u) { - const uint32_t nbN2 = (Nf + 63u) / 64u; - const uint32_t nbM2 = (m + 63u) / 64u; - gr.dispatch_at(fused_idx).workgroup_count_x = nbN2 * nbM2; - gr.dispatch_at(sep0).workgroup_count_x = 0u; - gr.dispatch_at(sep1).workgroup_count_x = 0u; - gr.dispatch_at(sep2).workgroup_count_x = 0u; - } else { - gr.dispatch_at(fused_idx).workgroup_count_x = 0u; - } - }); -} - void WebGPUGraph::copy_inputs(const std::vector& inputs) { for (size_t i = 0; i < inputs.size() && i < input_ids_.size(); i++) { const InputData& in = inputs[i]; @@ -1671,10 +1436,15 @@ void WebGPUGraph::copy_inputs(const std::vector& inputs) { // Upload only the live (cur) bytes, not the max allocation; cur_nbytes == // nbytes on a static graph, so this is byte-identical there. const size_t live_nbytes = tensor.cur_nbytes; + const bool buffer_is_fp16 = !tensor.is_int && tensor.elem_size == 2; + if (buffer_is_fp16 && !in.host_is_fp32) { + throw std::runtime_error( + "WebGPU: fp16 device input requires an fp32 host tensor"); + } // Fast path: host and GPU element types match byte-for-byte. if (in.nbytes == live_nbytes) { - wgpuQueueWriteBuffer(queue_, tensor.buffer, 0, in.data, live_nbytes); + write_storage_buffer(queue_, tensor.buffer, in.data, live_nbytes); continue; } @@ -1694,8 +1464,21 @@ void WebGPUGraph::copy_inputs(const std::vector& inputs) { #endif narrowed[e] = static_cast(src[e]); } - wgpuQueueWriteBuffer( - queue_, tensor.buffer, 0, narrowed.data(), live_nbytes); + write_storage_buffer(queue_, tensor.buffer, narrowed.data(), live_nbytes); + continue; + } + + // Require an explicit fp32 host dtype, not merely "not int64": inferring + // the narrow from the 2:1 byte ratio alone would silently reinterpret a + // same-sized non-fp32 host buffer (e.g. a stale int32) as fp32. + if (in.host_is_fp32 && buffer_is_fp16 && in.nbytes == live_nbytes * 2) { + const size_t numel = live_nbytes / sizeof(uint16_t); + const float* src = static_cast(in.data); + std::vector narrowed(numel); + for (size_t e = 0; e < numel; e++) { + narrowed[e] = executorch::runtime::etensor::Half(src[e]); + } + write_storage_buffer(queue_, tensor.buffer, narrowed.data(), live_nbytes); continue; } @@ -1707,36 +1490,162 @@ void WebGPUGraph::copy_inputs(const std::vector& inputs) { } } +#ifdef WGPU_BACKEND_ENABLE_PROFILING +// Profiling/attestation only; never compiled into a production build. Written +// during WebGPUGraph::execute without synchronization: the attestation +// harnesses that read them run one graph on one thread, so no atomics or +// locking are needed. To support concurrent profiled execution, make these +// per-instance behind a whole-record mutex (per-field atomics would not cover +// the conflict check's read-modify-write across both globals). +uint32_t g_last_route_mask = 0; +uint32_t g_last_route_conflict_count = 0; +#endif // WGPU_BACKEND_ENABLE_PROFILING + namespace { +#ifdef WGPU_BACKEND_ENABLE_PROFILING +constexpr uint32_t kRoutePrefill = 1u << 0; +constexpr uint32_t kRouteK16 = 1u << 1; +constexpr uint32_t kRouteMaterializedAttention = 1u << 2; +constexpr uint32_t kRouteT0Steel = 1u << 3; +constexpr uint32_t kRouteT1Bk64 = 1u << 4; +constexpr uint32_t kRouteT1Bk64Qkv = 1u << 5; +constexpr uint32_t kRouteT2PairedGateUp = 1u << 6; +constexpr uint32_t kRouteFusedSwiGlu = 1u << 7; +constexpr uint32_t kRouteGenericFallback = 1u << 8; +// bit 1u << 9 is intentionally reserved (a retired route) and left unused. +constexpr uint32_t kRouteFlashDecoding = 1u << 10; +constexpr uint32_t kRouteK16CausalBound = 1u << 11; +constexpr uint32_t kRouteBicolSubgroup = 1u << 12; +constexpr uint32_t kRouteQwen3Q16K16 = 1u << 13; +constexpr uint32_t kRouteQwen3Q32K16 = 1u << 14; +#endif // WGPU_BACKEND_ENABLE_PROFILING + // Bench gate: compiled out unless WGPU_BACKEND_ENABLE_PROFILING; then the // WEBGPU_TIMESTAMP_QUERY env var enables per-pass GPU timestamp queries. bool should_timestamp_query() { #ifdef WGPU_BACKEND_ENABLE_PROFILING - static const bool enabled = std::getenv("WEBGPU_TIMESTAMP_QUERY") != nullptr; - return enabled; + return std::getenv("WEBGPU_TIMESTAMP_QUERY") != nullptr; #else return false; #endif } } // namespace -void WebGPUGraph::execute() { +#ifdef WGPU_BACKEND_ENABLE_PROFILING +void WebGPUGraph::record_active_route(const std::string& kernel_name) { + uint32_t bits = 0; + if (kernel_name == "sdpa_streaming_attention_qwen3_q32_k16_causal_bound") { + bits = kRoutePrefill | kRouteK16CausalBound | kRouteQwen3Q32K16; + } else if (kernel_name == "sdpa_streaming_attention_qwen3_k16_causal_bound") { + bits = kRoutePrefill | kRouteK16CausalBound | kRouteQwen3Q16K16; + } else if ( + kernel_name.rfind("sdpa_streaming_attention_", 0) == 0 && + kernel_name.find("k16_causal_bound") != std::string::npos) { + bits = kRoutePrefill | kRouteK16CausalBound; + } else if (kernel_name == "sdpa_streaming_attention_k16") { + bits = kRoutePrefill | kRouteK16; + } else if ( + kernel_name.rfind("sdpa_compute_", 0) == 0 || + kernel_name == "sdpa_softmax") { + bits = kRoutePrefill | kRouteMaterializedAttention; + } else if (kernel_name == "fd_split" || kernel_name == "fd_reduce") { + bits = kRouteFlashDecoding; + } else if (kernel_name == "linear_q4gsw_coop4_bicol_subgroup") { + bits = kRouteBicolSubgroup; + } else if (kernel_name.rfind("linear_q4gsw_bk64_qkv", 0) == 0) { + bits = kRouteT1Bk64Qkv; + } else if (kernel_name.rfind("linear_q4gsw_bk64", 0) == 0) { + bits = kRouteT1Bk64; + } else if (kernel_name.rfind("linear_q4gsw_paired_gate_up", 0) == 0) { + bits = kRouteT2PairedGateUp; + } else if (kernel_name == "silu_mul_fused") { + bits = kRouteFusedSwiGlu; + } else if (kernel_name.rfind("linear_q4gsw_steel", 0) == 0) { + bits = kRouteT0Steel; + } else if (kernel_name.rfind("linear_q4gsw", 0) == 0) { + bits = kRouteGenericFallback; + } + + constexpr uint32_t kAttentionRoutes = kRouteK16 | kRouteK16CausalBound | + kRouteMaterializedAttention | kRouteFlashDecoding; + const uint32_t new_attention = bits & kAttentionRoutes; + const uint32_t prior_attention = g_last_route_mask & kAttentionRoutes; + if (new_attention != 0 && prior_attention != 0 && + (new_attention & prior_attention) == 0) { + ++g_last_route_conflict_count; + } + g_last_route_mask |= bits; +} +#endif // WGPU_BACKEND_ENABLE_PROFILING + +WebGPUExecutionPlan WebGPUGraph::make_execution_plan( + const WebGPUGraphExecutionOptions& options) const { + const size_t n = dispatches_.size(); + std::vector enabled_dispatches(n, true); + for (size_t i = 0; i < n; i++) { + if (dispatches_[i].kind != WebGPUDispatch::Kind::Compute) { + continue; + } + const bool zero_x = dispatches_[i].workgroup_count_x == 0; + const bool zero_y = dispatches_[i].workgroup_count_y == 0; + if (zero_x != zero_y) { + throw std::runtime_error("WebGPU: dispatch has a half-zero grid"); + } + enabled_dispatches[i] = !zero_x; + } + return plan_webgpu_execution( + n, + output_copies_.size(), + execute_config_, + suppressible_outputs_, + options, + enabled_dispatches); +} + +size_t WebGPUGraph::execute(const WebGPUExecutionPlan& plan) { +#ifdef WGPU_BACKEND_ENABLE_PROFILING + g_last_route_mask = 0; + g_last_route_conflict_count = 0; +#endif // WGPU_BACKEND_ENABLE_PROFILING const size_t n = dispatches_.size(); const size_t chunk = execute_config_.chunk_size; + if (plan.copy_outputs.size() != output_copies_.size()) { + throw std::runtime_error("WebGPU: execution plan output count mismatch"); + } + for (const auto& dispatch_chunk : plan.dispatch_chunks) { + for (size_t dispatch_index : dispatch_chunk) { + if (dispatch_index >= n) { + throw std::runtime_error( + "WebGPU: execution plan dispatch index out of range"); + } + } + } + + if (plan.dispatch_chunks.empty()) { + return 0; + } if (chunk == 0 || n <= chunk) { #ifdef WGPU_BACKEND_ENABLE_PROFILING + size_t active_compute_count = 0; + for (size_t i : plan.dispatch_chunks.front()) { + if (dispatches_[i].kind == WebGPUDispatch::Kind::Compute) { + active_compute_count++; + } + } // Bench: timestamp-query pool, null unless env-gated + feature present. WebGPUQueryPool* qp = nullptr; - if (should_timestamp_query() && n > 0) { + if (should_timestamp_query() && active_compute_count > 0) { if (auto* ctx = get_default_webgpu_context()) { if (ctx->timestamp_supported) { - if (!ctx->querypool || ctx->querypool->capacity() < n) { + if (!ctx->querypool || + ctx->querypool->capacity() < active_compute_count) { ctx->querypool = std::make_unique(); - ctx->querypool->initialize(device_, static_cast(n)); + ctx->querypool->initialize( + device_, static_cast(active_compute_count)); } qp = ctx->querypool.get(); - qp->reset(static_cast(n)); + qp->reset(static_cast(active_compute_count)); } } } @@ -1747,50 +1656,65 @@ void WebGPUGraph::execute() { wgpuDeviceCreateCommandEncoder(device_, &enc_desc); // One pass per dispatch: enforces storage RAW ordering across deps. - for (size_t i = 0; i < n; i++) { - const auto& dispatch = dispatches_[i]; - if (dispatch.kind == WebGPUDispatch::Kind::Copy) { - wgpuCommandEncoderCopyBufferToBuffer( - encoder, - dispatch.copy_src, - 0, - dispatch.copy_dst, - 0, - dispatch.copy_nbytes); - continue; - } - WGPUComputePassDescriptor pass_desc = {}; #ifdef WGPU_BACKEND_ENABLE_PROFILING - // tw must outlive BeginComputePass (the descriptor points at it). - WGPUPassTimestampWrites tw = {}; - if (qp) { - tw = qp->writes_for(static_cast(i)); - pass_desc.timestampWrites = &tw; - } + uint32_t query_index = 0; +#endif + for (const auto& dispatch_chunk : plan.dispatch_chunks) { + for (size_t i : dispatch_chunk) { + const auto& dispatch = dispatches_[i]; + if (dispatch.kind == WebGPUDispatch::Kind::Copy) { + wgpuCommandEncoderCopyBufferToBuffer( + encoder, + dispatch.copy_src, + 0, + dispatch.copy_dst, + 0, + dispatch.copy_nbytes); + continue; + } +#ifdef WGPU_BACKEND_ENABLE_PROFILING + record_active_route(dispatch.kernel_name); #endif // WGPU_BACKEND_ENABLE_PROFILING - WGPUComputePassEncoder pass = - wgpuCommandEncoderBeginComputePass(encoder, &pass_desc); - wgpuComputePassEncoderSetPipeline(pass, dispatch.pipeline); - wgpuComputePassEncoderSetBindGroup( - pass, 0, dispatch.bind_group, 0, nullptr); - wgpuComputePassEncoderDispatchWorkgroups( - pass, dispatch.workgroup_count_x, dispatch.workgroup_count_y, 1); - wgpuComputePassEncoderEnd(pass); - wgpuComputePassEncoderRelease(pass); + WGPUComputePassDescriptor pass_desc = {}; #ifdef WGPU_BACKEND_ENABLE_PROFILING - if (qp) { - qp->record( - static_cast(i), - dispatch.kernel_name, - {dispatch.workgroup_count_x, dispatch.workgroup_count_y, 1}, - {1, 1, 1}); - } + // tw must outlive BeginComputePass (the descriptor points at it). + WGPUPassTimestampWrites tw = {}; + if (qp) { + tw = qp->writes_for(query_index); + pass_desc.timestampWrites = &tw; + } #endif // WGPU_BACKEND_ENABLE_PROFILING + WGPUComputePassEncoder pass = + wgpuCommandEncoderBeginComputePass(encoder, &pass_desc); + wgpuComputePassEncoderSetPipeline(pass, dispatch.pipeline); + wgpuComputePassEncoderSetBindGroup( + pass, 0, dispatch.bind_group, 0, nullptr); + wgpuComputePassEncoderDispatchWorkgroups( + pass, dispatch.workgroup_count_x, dispatch.workgroup_count_y, 1); + wgpuComputePassEncoderEnd(pass); + wgpuComputePassEncoderRelease(pass); +#ifdef WGPU_BACKEND_ENABLE_PROFILING + if (qp) { + qp->record( + query_index, + dispatch.kernel_name, + {dispatch.workgroup_count_x, dispatch.workgroup_count_y, 1}, + {1, 1, 1}); + query_index++; + } +#endif // WGPU_BACKEND_ENABLE_PROFILING + } } - for (const auto& copy : output_copies_) { + for (size_t i = 0; i < output_copies_.size(); i++) { + const size_t logical_nbytes = tensors_[output_ids_[i]].cur_nbytes; + if (!plan.copy_outputs[i] || logical_nbytes == 0) { + continue; + } + const size_t copy_nbytes = storage_buffer_size(logical_nbytes); + const auto& copy = output_copies_[i]; wgpuCommandEncoderCopyBufferToBuffer( - encoder, copy.src_buffer, 0, copy.staging_buffer, 0, copy.nbytes); + encoder, copy.src_buffer, 0, copy.staging_buffer, 0, copy_nbytes); } #ifdef WGPU_BACKEND_ENABLE_PROFILING @@ -1812,7 +1736,7 @@ void WebGPUGraph::execute() { qp->print_results(); } #endif // WGPU_BACKEND_ENABLE_PROFILING - return; + return 1; } // GPU timestamp queries assume one submit; chunked execute is multi-submit. @@ -1822,21 +1746,13 @@ void WebGPUGraph::execute() { "(multi-submit); disable chunking to use GPU timestamp queries"); } - const size_t first_chunk = execute_config_.initial_chunk_size > 0 - ? execute_config_.initial_chunk_size - : chunk; - - size_t start = 0; - size_t current_chunk = first_chunk; - - while (start < n) { - size_t end = std::min(start + current_chunk, n); - + for (size_t chunk_index = 0; chunk_index < plan.dispatch_chunks.size(); + chunk_index++) { WGPUCommandEncoderDescriptor enc_desc = {}; WGPUCommandEncoder encoder = wgpuDeviceCreateCommandEncoder(device_, &enc_desc); - for (size_t i = start; i < end; i++) { + for (size_t i : plan.dispatch_chunks[chunk_index]) { if (dispatches_[i].kind == WebGPUDispatch::Kind::Copy) { wgpuCommandEncoderCopyBufferToBuffer( encoder, @@ -1847,6 +1763,9 @@ void WebGPUGraph::execute() { dispatches_[i].copy_nbytes); continue; } +#ifdef WGPU_BACKEND_ENABLE_PROFILING + record_active_route(dispatches_[i].kernel_name); +#endif // WGPU_BACKEND_ENABLE_PROFILING WGPUComputePassDescriptor pass_desc = {}; WGPUComputePassEncoder pass = wgpuCommandEncoderBeginComputePass(encoder, &pass_desc); @@ -1862,10 +1781,16 @@ void WebGPUGraph::execute() { wgpuComputePassEncoderRelease(pass); } - if (end == n) { - for (const auto& copy : output_copies_) { + if (chunk_index + 1 == plan.dispatch_chunks.size()) { + for (size_t i = 0; i < output_copies_.size(); i++) { + const size_t logical_nbytes = tensors_[output_ids_[i]].cur_nbytes; + if (!plan.copy_outputs[i] || logical_nbytes == 0) { + continue; + } + const size_t copy_nbytes = storage_buffer_size(logical_nbytes); + const auto& copy = output_copies_[i]; wgpuCommandEncoderCopyBufferToBuffer( - encoder, copy.src_buffer, 0, copy.staging_buffer, 0, copy.nbytes); + encoder, copy.src_buffer, 0, copy.staging_buffer, 0, copy_nbytes); } } @@ -1875,10 +1800,8 @@ void WebGPUGraph::execute() { wgpuCommandBufferRelease(cmd); wgpuCommandEncoderRelease(encoder); - - start = end; - current_chunk = chunk; } + return plan.dispatch_chunks.size(); } namespace { @@ -1887,95 +1810,125 @@ struct MapCallbackData { WGPUMapAsyncStatus status = WGPUMapAsyncStatus_Error; }; +using MapCallbackDataPtr = std::shared_ptr; + void buffer_map_callback( WGPUMapAsyncStatus status, WGPUStringView /*message*/, void* userdata1, void* /*userdata2*/) { - auto* data = static_cast(userdata1); - data->status = status; + std::unique_ptr data_owner( + static_cast(userdata1)); + (*data_owner)->status = status; } } // namespace -void WebGPUGraph::copy_outputs(std::vector>& outputs) { +void WebGPUGraph::copy_outputs( + std::vector& outputs, + const WebGPUExecutionPlan& plan) { + if (plan.copy_outputs.size() != output_copies_.size()) { + throw std::runtime_error("WebGPU: execution plan output count mismatch"); + } const size_t count = std::min(outputs.size(), output_staging_buffers_.size()); - std::vector cb_data(count); - std::vector map_futures(count, WGPUFuture{}); - // Map each output's LIVE staging size (an int64 output is int32-backed). - std::vector map_nbytes(count, 0); + // Reject all dtype/size mismatches before issuing an asynchronous map. + for (size_t i = 0; i < count; i++) { + if (!plan.copy_outputs[i] || outputs[i].nbytes == 0) { + continue; + } + const auto& tensor = tensors_[output_ids_[i]]; + const size_t logical_nbytes = tensor.cur_nbytes; + if (logical_nbytes == 0) { + continue; + } + const size_t dst_nbytes = outputs[i].nbytes; + const bool is_double_width = + dst_nbytes % 2 == 0 && dst_nbytes / 2 == logical_nbytes; + const bool widen_fp16 = + is_double_width && !tensor.is_int && tensor.elem_size == 2; + const bool widen_int32 = + is_double_width && tensor.is_int && tensor.elem_size == 4; + const bool buffer_is_fp16 = !tensor.is_int && tensor.elem_size == 2; + if (buffer_is_fp16 && !outputs[i].host_is_fp32) { + throw std::runtime_error( + "WebGPU: fp16 device output requires an fp32 host tensor"); + } + if (outputs[i].host_is_fp32 && buffer_is_fp16 && !widen_fp16) { + throw std::runtime_error("WebGPU: fp16 output buffer size mismatch"); + } + if (dst_nbytes != logical_nbytes && !widen_fp16 && !widen_int32) { + throw std::runtime_error("WebGPU: output buffer size mismatch"); + } + } for (size_t i = 0; i < count; i++) { - map_nbytes[i] = tensors_[output_ids_[i]].cur_nbytes; - if (map_nbytes[i] == 0) { - cb_data[i].status = WGPUMapAsyncStatus_Success; + if (!plan.copy_outputs[i] || outputs[i].nbytes == 0) { continue; } + const auto& tensor = tensors_[output_ids_[i]]; + const size_t logical_nbytes = tensor.cur_nbytes; + if (logical_nbytes == 0) { + continue; + } + const size_t map_nbytes = storage_buffer_size(logical_nbytes); + const size_t dst_nbytes = outputs[i].nbytes; + const bool is_double_width = + dst_nbytes % 2 == 0 && dst_nbytes / 2 == logical_nbytes; + const bool widen_fp16 = + is_double_width && !tensor.is_int && tensor.elem_size == 2; + const bool widen_int32 = + is_double_width && tensor.is_int && tensor.elem_size == 4; + + const auto cb_data = std::make_shared(); WGPUBufferMapCallbackInfo cb_info = {}; cb_info.mode = WGPUCallbackMode_WaitAnyOnly; cb_info.callback = buffer_map_callback; - cb_info.userdata1 = &cb_data[i]; - map_futures[i] = wgpuBufferMapAsync( - output_staging_buffers_[i], - WGPUMapMode_Read, - 0, - map_nbytes[i], - cb_info); - } - - // Tracks which output buffers are currently mapped so a mid-loop throw can - // release them before propagating (no dangling mapped buffers). - std::vector is_mapped(count, false); - - try { - for (size_t i = 0; i < count; i++) { - if (map_nbytes[i] == 0) { - continue; - } - if (webgpu_wait(instance_, map_futures[i]) != WGPUWaitStatus_Success) { - throw std::runtime_error("WebGPU: WaitAny failed for output map"); - } - if (cb_data[i].status == WGPUMapAsyncStatus_Success) { - is_mapped[i] = true; + cb_info.userdata1 = new MapCallbackDataPtr(cb_data); + const WGPUFuture map_future = wgpuBufferMapAsync( + output_staging_buffers_[i], WGPUMapMode_Read, 0, map_nbytes, cb_info); + if (webgpu_wait(instance_, map_future) != WGPUWaitStatus_Success) { + // Cancel the outstanding request, then drain its WaitAny-only callback + // when possible. The callback owns a shared reference, so even a failed + // drain cannot leave it pointing at stack-owned storage. + // An undrained callback may still fire; leaking beats a use-after-free. + wgpuBufferUnmap(output_staging_buffers_[i]); + const WGPUWaitStatus drain_status = webgpu_wait(instance_, map_future); + if (drain_status != WGPUWaitStatus_Success) { + throw std::runtime_error( + "WebGPU: output map cancellation callback did not drain"); } + throw std::runtime_error("WebGPU: WaitAny failed for output map"); } - - for (size_t i = 0; i < count; i++) { - if (map_nbytes[i] == 0) { - continue; - } - if (cb_data[i].status != WGPUMapAsyncStatus_Success) { - throw std::runtime_error("WebGPU buffer map failed for output"); - } - const void* mapped = wgpuBufferGetConstMappedRange( - output_staging_buffers_[i], 0, map_nbytes[i]); - const size_t dst_nbytes = outputs[i].second; - if (dst_nbytes == map_nbytes[i]) { - std::memcpy(outputs[i].first, mapped, map_nbytes[i]); - } else if ( - dst_nbytes == 2 * map_nbytes[i] && tensors_[output_ids_[i]].is_int && - tensors_[output_ids_[i]].elem_size == 4) { - // int64 host output backed by an int32 GPU buffer: widen (sign-extend). - const int32_t* src = static_cast(mapped); - int64_t* dst = static_cast(outputs[i].first); - const size_t n = map_nbytes[i] / sizeof(int32_t); - for (size_t k = 0; k < n; k++) { - dst[k] = static_cast(src[k]); - } - } else { - throw std::runtime_error("WebGPU: output buffer size mismatch"); - } + if (cb_data->status != WGPUMapAsyncStatus_Success) { + throw std::runtime_error("WebGPU buffer map failed for output"); + } + const void* mapped = wgpuBufferGetConstMappedRange( + output_staging_buffers_[i], 0, map_nbytes); + if (mapped == nullptr) { wgpuBufferUnmap(output_staging_buffers_[i]); - is_mapped[i] = false; + throw std::runtime_error("WebGPU mapped output range is null"); } - } catch (...) { - for (size_t j = 0; j < count; j++) { - if (is_mapped[j]) { - wgpuBufferUnmap(output_staging_buffers_[j]); + if (widen_fp16) { + const auto* src = + static_cast(mapped); + auto* dst = static_cast(outputs[i].data); + const size_t n = logical_nbytes / sizeof(*src); + for (size_t k = 0; k < n; k++) { + dst[k] = static_cast(src[k]); } + } else if (widen_int32) { + // int64 host output backed by an int32 GPU buffer: widen (sign-extend). + const int32_t* src = static_cast(mapped); + int64_t* dst = static_cast(outputs[i].data); + const size_t n = logical_nbytes / sizeof(int32_t); + for (size_t k = 0; k < n; k++) { + dst[k] = static_cast(src[k]); + } + } else { + std::memcpy(outputs[i].data, mapped, logical_nbytes); } - throw; + wgpuBufferUnmap(output_staging_buffers_[i]); } } diff --git a/backends/webgpu/runtime/WebGPUGraph.h b/backends/webgpu/runtime/WebGPUGraph.h index daa083e992c..23ce9df03ed 100644 --- a/backends/webgpu/runtime/WebGPUGraph.h +++ b/backends/webgpu/runtime/WebGPUGraph.h @@ -12,11 +12,16 @@ #include #include +#include #include +#include #include #include +#include #include +#include +#include #include namespace executorch::backends::webgpu { @@ -34,6 +39,8 @@ struct WebGPUTensor { // Serialized (GPU-side) element type, used to narrow wider host inputs. size_t elem_size = 0; bool is_int = false; + // Exact BOOL tag for byte-packed WGSL storage. + bool is_bool = false; // Exactly int8 (not uint8/bool), so int8-only ops can guard their dtype. bool is_int8 = false; }; @@ -43,6 +50,15 @@ struct InputData { const void* data = nullptr; size_t nbytes = 0; bool host_is_int64 = false; + bool host_is_fp32 = false; +}; + +// Host destination for a graph output. host_is_fp32 gates the fp16->fp32 widen +// on readback (mirrors InputData's guard on the copy_inputs narrow path). +struct OutputData { + void* data = nullptr; + size_t nbytes = 0; + bool host_is_fp32 = false; }; struct WebGPUDispatch { @@ -59,6 +75,37 @@ struct WebGPUDispatch { size_t copy_nbytes = 0; }; +struct WebGPUBufferBinding { + WGPUBuffer buffer = nullptr; + uint64_t offset = 0; + uint64_t size = 0; +}; + +struct WebGPUSpecializationConstant { + std::string name; + double value = 0.0; +}; + +struct WebGPUDispatchGrid { + uint32_t x = 1; + uint32_t y = 1; +}; + +struct WebGPUComputeDispatchDescriptor { + std::string shader_name; + std::string entry_point = "main"; + std::string kernel_name; + std::vector bindings; + std::vector constants; + WebGPUDispatchGrid grid; +}; + +std::string make_compute_pipeline_key( + const WebGPUComputeDispatchDescriptor& descriptor); + +void validate_compute_dispatch_descriptor( + const WebGPUComputeDispatchDescriptor& descriptor); + struct OutputCopy { WGPUBuffer src_buffer = nullptr; WGPUBuffer staging_buffer = nullptr; @@ -74,11 +121,6 @@ struct ConstantSource { size_t nbytes = 0; }; -struct ExecuteConfig { - size_t chunk_size = 0; - size_t initial_chunk_size = 0; -}; - struct WebGPUMemoryStats { size_t tensor_buffer_bytes = 0; size_t shared_buffer_bytes = 0; @@ -96,6 +138,13 @@ struct WebGPUMemoryStats { } }; +struct WebGPUGraphConfig { + bool f16_kv_cache = false; + bool f16_accumulate_gemm = false; + int sdpa_query_tile = 0; + bool record_q4gsw_decode_route = false; +}; + class WebGPUGraph { public: WebGPUGraph(); @@ -106,19 +155,25 @@ class WebGPUGraph { void build( const void* flatbuffer_data, const uint8_t* constant_data, + size_t constant_data_size, const executorch::runtime::NamedDataMap* named_data_map = nullptr, - bool f16_kv_cache = false, - bool f16_accumulate_gemm = false); + WebGPUGraphConfig config = {}); // Copy input tensor data from host pointers into GPU buffers. void copy_inputs(const std::vector& inputs); - // Execute all recorded dispatches. - void execute(); + WebGPUExecutionPlan make_execution_plan( + const WebGPUGraphExecutionOptions& options) const; + + // Execute the dispatches selected by a plan created for this graph. Returns + // the number of GPU queue submissions performed. + size_t execute(const WebGPUExecutionPlan& plan); // Copy output tensor data from GPU buffers back to host pointers. // Uses mapAsync + ASYNCIFY in Wasm. - void copy_outputs(std::vector>& outputs); + void copy_outputs( + std::vector& outputs, + const WebGPUExecutionPlan& plan); const std::vector& input_ids() const { return input_ids_; @@ -197,14 +252,43 @@ class WebGPUGraph { symint_dim_sources_.push_back({symint_id, tensor_id, dim}); } + bool tensor_has_dynamic_dims(int tensor_id) const { + return dynamic_tensor_ids_.count(tensor_id) != 0; + } + + bool has_dynamic_shapes() const { + return !dynamic_tensor_ids_.empty(); + } + // Execute-time select_as_symint read; mirrors Vulkan select_as_symint_impl. void update_symints_from_inputs(const std::vector& inputs); // Per-SymInt resize hook; mirrors Vulkan DynamicDispatchNode::trigger_resize. void add_resize_hook(int symint_id, std::function fn) { + if (symint_id < 0 || symint_id >= num_values() || + get_value_type(symint_id) != ValueType::SymInt) { + throw std::runtime_error("WebGPU resize: trigger must be a SymInt"); + } + if (!fn) { + throw std::runtime_error("WebGPU resize: null SymInt resize hook"); + } resize_hooks_.push_back({symint_id, std::move(fn)}); } + template + void add_resize_hook( + int symint_id, + void (*fn)(WebGPUGraph&, const Context&), + Context context) { + if (fn == nullptr) { + throw std::runtime_error("WebGPU resize: null SymInt resize hook"); + } + add_resize_hook( + symint_id, [fn, context = std::move(context)](WebGPUGraph& graph) { + fn(graph, context); + }); + } + // Set a graph input's live dims (<= max) + dirty it; static path stays inert. void resize_input(int value_id, const std::vector& new_dims); // Set a tensor's live dims (an op resize hook calls this for its output to @@ -219,9 +303,31 @@ class WebGPUGraph { void add_tensor_resize_hook( int trigger_tensor_id, std::function fn) { + if (trigger_tensor_id < 0 || trigger_tensor_id >= num_values() || + get_value_type(trigger_tensor_id) != ValueType::Tensor) { + throw std::runtime_error("WebGPU resize: trigger must be a Tensor"); + } + if (!fn) { + throw std::runtime_error("WebGPU resize: null tensor resize hook"); + } tensor_resize_hooks_.push_back({trigger_tensor_id, std::move(fn)}); } + template + void add_tensor_resize_hook( + int trigger_tensor_id, + void (*fn)(WebGPUGraph&, const Context&), + Context context) { + if (fn == nullptr) { + throw std::runtime_error("WebGPU resize: null tensor resize hook"); + } + add_tensor_resize_hook( + trigger_tensor_id, + [fn, context = std::move(context)](WebGPUGraph& graph) { + fn(graph, context); + }); + } + // Run hooks for changed SymInts and tensors, then clear; call before execute. void propagate_resize(); @@ -233,6 +339,26 @@ class WebGPUGraph { return dispatches_.size(); } + size_t register_dispatch_route_group( + const std::vector& ranges) { + validate_dynamic_dispatch_route_ranges(ranges); + return dispatch_routes_.register_group( + dispatches_.size(), ranges, [&](size_t i) { + return dispatches_[i].kind == WebGPUDispatch::Kind::Compute; + }); + } + + void select_dispatch_route( + size_t group, + size_t active_route, + const std::vector& active_grids) { + dispatch_routes_.select( + group, active_route, active_grids, [&](size_t i, utils::WgCount grid) { + dispatches_[i].workgroup_count_x = grid.x; + dispatches_[i].workgroup_count_y = grid.y; + }); + } + WGPUDevice device() const { return device_; } @@ -286,6 +412,24 @@ class WebGPUGraph { owned_uniform_buffers_.push_back(buffer); } + template + WGPUBuffer create_params_buffer(const Block& data) { + static_assert( + std::is_trivially_copyable::value, + "WebGPU parameter blocks must be trivially copyable"); + static_assert( + sizeof(Block) % 4u == 0u, + "WebGPU parameter blocks must have a 4-byte-aligned size"); + WGPUBuffer buffer = make_uniform_buffer(&data, sizeof(Block)); + try { + own_uniform_buffer(buffer); + } catch (...) { + wgpuBufferRelease(buffer); + throw; + } + return buffer; + } + // Graph-owned scratch storage buffer for fused-op intermediates (e.g. SDPA). WGPUBuffer create_scratch_buffer(size_t nbytes); @@ -322,6 +466,26 @@ class WebGPUGraph { // in the memory stats. Shared helper for ops needing a uniform Params buffer. WGPUBuffer make_uniform_buffer(const void* data, size_t size); + size_t add_compute_dispatch( + const WebGPUComputeDispatchDescriptor& descriptor); + + template + size_t add_dynamic_compute_dispatch( + const WebGPUComputeDispatchDescriptor& descriptor, + int trigger_tensor_id, + WebGPUDispatchGrid (*pick_grid)(const WebGPUGraph&, const Context&), + Context context) { + if (pick_grid == nullptr) { + throw std::runtime_error("WebGPU dynamic dispatch: null grid picker"); + } + return add_dynamic_compute_dispatch_impl( + descriptor, + trigger_tensor_id, + [pick_grid, context = std::move(context)](const WebGPUGraph& graph) { + return pick_grid(graph, context); + }); + } + WGPUShaderModule get_or_create_shader( const std::string& key, const char* wgsl_source); @@ -365,6 +529,19 @@ class WebGPUGraph { return value_types_[id]; } + // Memory-aliasing group id for the tensor's shared buffer, or -1 if it has + // none; fusion passes use this to reject candidates aliased with something + // the planner may reuse outside the fusion's control. + int mem_obj_id(int id) const { + return tensor_mem_obj_ids_[id]; + } + + // True if id is a prepack-routed constant with a recorded source (inline + // offset or named-data-map key); fusion passes require direct constants. + bool has_constant_source(int id) const { + return constant_sources_.count(id) != 0; + } + public: // True when the sdpa K/V cache is stored f16-packed (runtime opt-in). bool kv_f16() const { @@ -374,13 +551,27 @@ class WebGPUGraph { // True when the q4gsw steel prefill GEMM uses the lossy f16-accumulate kernel // (runtime opt-in; perplexity-gated, not bit-exact). bool f16_accumulate_gemm() const { - return f16_accumulate_gemm_; + return config_.f16_accumulate_gemm; + } + + // Runtime-selected SDPA query-tile candidate; 0 = geometry default (Q16), + // 32 = Q32 candidate. + int sdpa_query_tile() const { + return config_.sdpa_query_tile; + } + + const WebGPUGraphConfig& config() const { + return config_; } private: +#ifdef WGPU_BACKEND_ENABLE_PROFILING + void record_active_route(const std::string& kernel_name); +#endif // WGPU_BACKEND_ENABLE_PROFILING + bool kv_f16_ = false; std::unordered_set kv_cache_ids_; - bool f16_accumulate_gemm_ = false; + WebGPUGraphConfig config_; private: WGPUInstance instance_ = nullptr; @@ -406,6 +597,7 @@ class WebGPUGraph { std::unordered_map symints_; std::vector symint_sources_; std::vector symint_dim_sources_; + std::unordered_set dynamic_tensor_ids_; // Resize hooks + the set of SymInts changed since the last propagate_resize. struct ResizeHook { @@ -424,6 +616,29 @@ class WebGPUGraph { std::vector tensor_resize_hooks_; std::unordered_set dirty_tensors_; + // Dynamic grids are stored separately so ordinary dispatches remain compact. + // The graph owns each dispatch index and picker; ops only provide typed + // context and a named grid function. + struct DynamicDispatchGrid { + size_t dispatch_index; + int trigger_tensor_id; + std::function pick_grid; + }; + std::vector dynamic_dispatch_grids_; + + struct PendingDynamicDispatchGrid { + size_t dispatch_index; + WebGPUDispatchGrid grid; + }; + std::vector pending_dynamic_dispatch_grids_; + + size_t add_dynamic_compute_dispatch_impl( + const WebGPUComputeDispatchDescriptor& descriptor, + int trigger_tensor_id, + std::function pick_grid); + void validate_dynamic_dispatch_route_ranges( + const std::vector& ranges) const; + std::vector input_ids_; std::vector output_ids_; @@ -454,12 +669,16 @@ class WebGPUGraph { // Pre-computed output copy descriptors for execute(). std::vector output_copies_; + std::vector suppressible_outputs_; + std::vector dispatches_; + utils::DispatchRouteRegistry dispatch_routes_; // Prepack-routed constant sources (offset/named-key + size); the prepack node // materializes these once. constant_data_/named_data_map_ point at the .pte // bytes and are valid only during build(). const uint8_t* constant_data_ = nullptr; + size_t constant_data_size_ = 0; const executorch::runtime::NamedDataMap* named_data_map_ = nullptr; std::unordered_map constant_sources_; @@ -471,43 +690,11 @@ class WebGPUGraph { std::unordered_map bgl_cache_; size_t uniform_buffer_bytes_ = 0; - - // QKV-concat fusion: one detected attention q/k/v linear - // triple sharing an input activation (value ids + shapes), fused in build() - // into a single multi-output q4gsw GEMM that scatter-writes q/k/v. Only used - // during build(); inert (never populated) when no q/k/v triple matches. - struct QkvFusionGroup { - int input_id = -1; - int out_q = -1, out_k = -1, out_v = -1; - int weight_q = -1, weight_k = -1, weight_v = -1; - int scales_q = -1, scales_k = -1, scales_v = -1; - uint32_t Nq = 0, Nk = 0, Nv = 0; // 2048, 512, 512 - uint32_t K = 0, K_packed = 0, group_size = 0, num_groups = 0; - uint32_t padded_N_q = 0, padded_N_k = 0, padded_N_v = 0; - unsigned op_idx[3] = {0, 0, 0}; // the 3 q/k/v linear op-chain indices - size_t sep_dispatch[3] = { - 0, - 0, - 0}; // their dispatch indices (filled in build()) - size_t fused_dispatch = 0; // the fused GEMM dispatch index - WGPUBuffer fused_params = - nullptr; // the fused params UBO (rewritten by the hook) - }; - // Concat the 3 packed weights (row-stack) + scales (strided gather) into - // fused buffers, then record ONE fused-GEMM dispatch (bespoke 8-binding - // layout) that writes the 3 original q/k/v output buffers, plus a 3-output - // resize hook. - void add_qkv_fused_dispatch(QkvFusionGroup& g); - void add_qkv_fused_hook(const QkvFusionGroup& g); - - // SwiGLU fusion: emit ONE fused elementwise dispatch - // computing out = (gate * sigmoid(gate)) * up, replacing the sigmoid + 2 - // muls. `out` is repointed to a private pooled buffer (aliasing guard); - // `gate` is likewise given a private pooled buffer at its producer op by the - // build() walk (the planner reuse-aliases up onto gate's slot, so up_proj - // would stomp gate before the fused reads it). Only used during build(); the - // detection maps are empty (inert) when no SwiGLU triple matches. - void add_swiglu_fused_dispatch(int gate_id, int up_id, int out_id); }; +#ifdef WGPU_BACKEND_ENABLE_PROFILING +extern uint32_t g_last_route_mask; +extern uint32_t g_last_route_conflict_count; +#endif // WGPU_BACKEND_ENABLE_PROFILING + } // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/WebGPUShaderRegistry.cpp b/backends/webgpu/runtime/WebGPUShaderRegistry.cpp new file mode 100644 index 00000000000..25de35d49d5 --- /dev/null +++ b/backends/webgpu/runtime/WebGPUShaderRegistry.cpp @@ -0,0 +1,1108 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +// @generated by scripts/gen_wgsl_headers.py - DO NOT EDIT. + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +namespace executorch::backends::webgpu { +namespace { + +constexpr std::array kShaderRegistry = {{ + { + "abs", + kAbsWGSL, + kAbsWorkgroupSizeX, + kAbsWorkgroupSizeY, + kAbsWorkgroupSizeZ, + }, + { + "adamw_step", + kAdamwStepWGSL, + kAdamwStepWorkgroupSizeX, + kAdamwStepWorkgroupSizeY, + kAdamwStepWorkgroupSizeZ, + }, + { + "addmm_tiled", + kAddmmTiledWGSL, + kAddmmTiledWorkgroupSizeX, + kAddmmTiledWorkgroupSizeY, + kAddmmTiledWorkgroupSizeZ, + }, + { + "amax", + kAmaxWGSL, + kAmaxWorkgroupSizeX, + kAmaxWorkgroupSizeY, + kAmaxWorkgroupSizeZ, + }, + { + "amin", + kAminWGSL, + kAminWorkgroupSizeX, + kAminWorkgroupSizeY, + kAminWorkgroupSizeZ, + }, + { + "apply_rotary_emb_interleaved", + kApplyRotaryEmbInterleavedWGSL, + kApplyRotaryEmbInterleavedWorkgroupSizeX, + kApplyRotaryEmbInterleavedWorkgroupSizeY, + kApplyRotaryEmbInterleavedWorkgroupSizeZ, + }, + { + "arg_reduce", + kArgReduceWGSL, + kArgReduceWorkgroupSizeX, + kArgReduceWorkgroupSizeY, + kArgReduceWorkgroupSizeZ, + }, + { + "avg_pool2d", + kAvgPool2dWGSL, + kAvgPool2dWorkgroupSizeX, + kAvgPool2dWorkgroupSizeY, + kAvgPool2dWorkgroupSizeZ, + }, + { + "batch_norm", + kBatchNormWGSL, + kBatchNormWorkgroupSizeX, + kBatchNormWorkgroupSizeY, + kBatchNormWorkgroupSizeZ, + }, + { + "binary_add", + kBinaryAddWGSL, + kBinaryAddWorkgroupSizeX, + kBinaryAddWorkgroupSizeY, + kBinaryAddWorkgroupSizeZ, + }, + { + "binary_div", + kBinaryDivWGSL, + kBinaryDivWorkgroupSizeX, + kBinaryDivWorkgroupSizeY, + kBinaryDivWorkgroupSizeZ, + }, + { + "binary_floor_divide", + kBinaryFloorDivideWGSL, + kBinaryFloorDivideWorkgroupSizeX, + kBinaryFloorDivideWorkgroupSizeY, + kBinaryFloorDivideWorkgroupSizeZ, + }, + { + "binary_minimum", + kBinaryMinimumWGSL, + kBinaryMinimumWorkgroupSizeX, + kBinaryMinimumWorkgroupSizeY, + kBinaryMinimumWorkgroupSizeZ, + }, + { + "binary_mul", + kBinaryMulWGSL, + kBinaryMulWorkgroupSizeX, + kBinaryMulWorkgroupSizeY, + kBinaryMulWorkgroupSizeZ, + }, + { + "binary_pow", + kBinaryPowWGSL, + kBinaryPowWorkgroupSizeX, + kBinaryPowWorkgroupSizeY, + kBinaryPowWorkgroupSizeZ, + }, + { + "binary_sub", + kBinarySubWGSL, + kBinarySubWorkgroupSizeX, + kBinarySubWorkgroupSizeY, + kBinarySubWorkgroupSizeZ, + }, + { + "bitwise_not", + kBitwiseNotWGSL, + kBitwiseNotWorkgroupSizeX, + kBitwiseNotWorkgroupSizeY, + kBitwiseNotWorkgroupSizeZ, + }, + { + "bmm_tiled", + kBmmTiledWGSL, + kBmmTiledWorkgroupSizeX, + kBmmTiledWorkgroupSizeY, + kBmmTiledWorkgroupSizeZ, + }, + { + "bmm_vec4", + kBmmVec4WGSL, + kBmmVec4WorkgroupSizeX, + kBmmVec4WorkgroupSizeY, + kBmmVec4WorkgroupSizeZ, + }, + { + "cat", + kCatWGSL, + kCatWorkgroupSizeX, + kCatWorkgroupSizeY, + kCatWorkgroupSizeZ, + }, + { + "choose_qparams_affine", + kChooseQparamsAffineWGSL, + kChooseQparamsAffineWorkgroupSizeX, + kChooseQparamsAffineWorkgroupSizeY, + kChooseQparamsAffineWorkgroupSizeZ, + }, + { + "clamp", + kClampWGSL, + kClampWorkgroupSizeX, + kClampWorkgroupSizeY, + kClampWorkgroupSizeZ, + }, + { + "compare", + kCompareWGSL, + kCompareWorkgroupSizeX, + kCompareWorkgroupSizeY, + kCompareWorkgroupSizeZ, + }, + { + "compare_eq", + kCompareEqWGSL, + kCompareEqWorkgroupSizeX, + kCompareEqWorkgroupSizeY, + kCompareEqWorkgroupSizeZ, + }, + { + "compare_ge", + kCompareGeWGSL, + kCompareGeWorkgroupSizeX, + kCompareGeWorkgroupSizeY, + kCompareGeWorkgroupSizeZ, + }, + { + "compare_gt", + kCompareGtWGSL, + kCompareGtWorkgroupSizeX, + kCompareGtWorkgroupSizeY, + kCompareGtWorkgroupSizeZ, + }, + { + "compare_le", + kCompareLeWGSL, + kCompareLeWorkgroupSizeX, + kCompareLeWorkgroupSizeY, + kCompareLeWorkgroupSizeZ, + }, + { + "compare_lt", + kCompareLtWGSL, + kCompareLtWorkgroupSizeX, + kCompareLtWorkgroupSizeY, + kCompareLtWorkgroupSizeZ, + }, + { + "compare_ne", + kCompareNeWGSL, + kCompareNeWorkgroupSizeX, + kCompareNeWorkgroupSizeY, + kCompareNeWorkgroupSizeZ, + }, + { + "constant_pad_nd", + kConstantPadNdWGSL, + kConstantPadNdWorkgroupSizeX, + kConstantPadNdWorkgroupSizeY, + kConstantPadNdWorkgroupSizeZ, + }, + { + "conv1d_dw", + kConv1dDwWGSL, + kConv1dDwWorkgroupSizeX, + kConv1dDwWorkgroupSizeY, + kConv1dDwWorkgroupSizeZ, + }, + { + "conv1d_pw", + kConv1dPwWGSL, + kConv1dPwWorkgroupSizeX, + kConv1dPwWorkgroupSizeY, + kConv1dPwWorkgroupSizeZ, + }, + { + "conv2d", + kConv2dWGSL, + kConv2dWorkgroupSizeX, + kConv2dWorkgroupSizeY, + kConv2dWorkgroupSizeZ, + }, + { + "conv2d_gemm", + kConv2dGemmWGSL, + kConv2dGemmWorkgroupSizeX, + kConv2dGemmWorkgroupSizeY, + kConv2dGemmWorkgroupSizeZ, + }, + { + "conv2d_vec4", + kConv2dVec4WGSL, + kConv2dVec4WorkgroupSizeX, + kConv2dVec4WorkgroupSizeY, + kConv2dVec4WorkgroupSizeZ, + }, + { + "conv_transpose2d", + kConvTranspose2dWGSL, + kConvTranspose2dWorkgroupSizeX, + kConvTranspose2dWorkgroupSizeY, + kConvTranspose2dWorkgroupSizeZ, + }, + { + "conv_with_clamp", + kConvWithClampWGSL, + kConvWithClampWorkgroupSizeX, + kConvWithClampWorkgroupSizeY, + kConvWithClampWorkgroupSizeZ, + }, + { + "cos", + kCosWGSL, + kCosWorkgroupSizeX, + kCosWorkgroupSizeY, + kCosWorkgroupSizeZ, + }, + { + "dequantize_per_tensor", + kDequantizePerTensorWGSL, + kDequantizePerTensorWorkgroupSizeX, + kDequantizePerTensorWorkgroupSizeY, + kDequantizePerTensorWorkgroupSizeZ, + }, + { + "embedding", + kEmbeddingWGSL, + kEmbeddingWorkgroupSizeX, + kEmbeddingWorkgroupSizeY, + kEmbeddingWorkgroupSizeZ, + }, + { + "embedding_q4gsw", + kEmbeddingQ4gswWGSL, + kEmbeddingQ4gswWorkgroupSizeX, + kEmbeddingQ4gswWorkgroupSizeY, + kEmbeddingQ4gswWorkgroupSizeZ, + }, + { + "et_vk_sdpa_av", + kEtVkSdpaAvWGSL, + kEtVkSdpaAvWorkgroupSizeX, + kEtVkSdpaAvWorkgroupSizeY, + kEtVkSdpaAvWorkgroupSizeZ, + }, + { + "et_vk_sdpa_qk", + kEtVkSdpaQkWGSL, + kEtVkSdpaQkWorkgroupSizeX, + kEtVkSdpaQkWorkgroupSizeY, + kEtVkSdpaQkWorkgroupSizeZ, + }, + { + "et_vk_sdpa_qk_entry", + kEtVkSdpaQkEntryWGSL, + kEtVkSdpaQkEntryWorkgroupSizeX, + kEtVkSdpaQkEntryWorkgroupSizeY, + kEtVkSdpaQkEntryWorkgroupSizeZ, + }, + { + "exp", + kExpWGSL, + kExpWorkgroupSizeX, + kExpWorkgroupSizeY, + kExpWorkgroupSizeZ, + }, + { + "expand_copy", + kExpandCopyWGSL, + kExpandCopyWorkgroupSizeX, + kExpandCopyWorkgroupSizeY, + kExpandCopyWorkgroupSizeZ, + }, + { + "fill", + kFillWGSL, + kFillWorkgroupSizeX, + kFillWorkgroupSizeY, + kFillWorkgroupSizeZ, + }, + { + "flip", + kFlipWGSL, + kFlipWorkgroupSizeX, + kFlipWorkgroupSizeY, + kFlipWorkgroupSizeZ, + }, + { + "fused_ce", + kFusedCeWGSL, + kFusedCeWorkgroupSizeX, + kFusedCeWorkgroupSizeY, + kFusedCeWorkgroupSizeZ, + }, + { + "gather", + kGatherWGSL, + kGatherWorkgroupSizeX, + kGatherWorkgroupSizeY, + kGatherWorkgroupSizeZ, + }, + { + "gelu", + kGeluWGSL, + kGeluWorkgroupSizeX, + kGeluWorkgroupSizeY, + kGeluWorkgroupSizeZ, + }, + { + "grid_priors", + kGridPriorsWGSL, + kGridPriorsWorkgroupSizeX, + kGridPriorsWorkgroupSizeY, + kGridPriorsWorkgroupSizeZ, + }, + { + "grid_sampler_2d", + kGridSampler2dWGSL, + kGridSampler2dWorkgroupSizeX, + kGridSampler2dWorkgroupSizeY, + kGridSampler2dWorkgroupSizeZ, + }, + { + "group_norm", + kGroupNormWGSL, + kGroupNormWorkgroupSizeX, + kGroupNormWorkgroupSizeY, + kGroupNormWorkgroupSizeZ, + }, + { + "group_norm_reduce", + kGroupNormReduceWGSL, + kGroupNormReduceWorkgroupSizeX, + kGroupNormReduceWorkgroupSizeY, + kGroupNormReduceWorkgroupSizeZ, + }, + { + "hardswish", + kHardswishWGSL, + kHardswishWorkgroupSizeX, + kHardswishWorkgroupSizeY, + kHardswishWorkgroupSizeZ, + }, + { + "index", + kIndexWGSL, + kIndexWorkgroupSizeX, + kIndexWorkgroupSizeY, + kIndexWorkgroupSizeZ, + }, + { + "index_select", + kIndexSelectWGSL, + kIndexSelectWorkgroupSizeX, + kIndexSelectWorkgroupSizeY, + kIndexSelectWorkgroupSizeZ, + }, + { + "leaky_relu", + kLeakyReluWGSL, + kLeakyReluWorkgroupSizeX, + kLeakyReluWorkgroupSizeY, + kLeakyReluWorkgroupSizeZ, + }, + { + "linear_dW", + kLinearDwWGSL, + kLinearDwWorkgroupSizeX, + kLinearDwWorkgroupSizeY, + kLinearDwWorkgroupSizeZ, + }, + { + "linear_dq8ca_q4gsw", + kLinearDq8caQ4gswWGSL, + kLinearDq8caQ4gswWorkgroupSizeX, + kLinearDq8caQ4gswWorkgroupSizeY, + kLinearDq8caQ4gswWorkgroupSizeZ, + }, + { + "linear_q8ta_q8csw", + kLinearQ8taQ8cswWGSL, + kLinearQ8taQ8cswWorkgroupSizeX, + kLinearQ8taQ8cswWorkgroupSizeY, + kLinearQ8taQ8cswWorkgroupSizeZ, + }, + { + "linear_tiled", + kLinearTiledWGSL, + kLinearTiledWorkgroupSizeX, + kLinearTiledWorkgroupSizeY, + kLinearTiledWorkgroupSizeZ, + }, + { + "linear_vec4", + kLinearVec4WGSL, + kLinearVec4WorkgroupSizeX, + kLinearVec4WorkgroupSizeY, + kLinearVec4WorkgroupSizeZ, + }, + { + "log_softmax", + kLogSoftmaxWGSL, + kLogSoftmaxWorkgroupSizeX, + kLogSoftmaxWorkgroupSizeY, + kLogSoftmaxWorkgroupSizeZ, + }, + { + "logical_and", + kLogicalAndWGSL, + kLogicalAndWorkgroupSizeX, + kLogicalAndWorkgroupSizeY, + kLogicalAndWorkgroupSizeZ, + }, + { + "logical_not", + kLogicalNotWGSL, + kLogicalNotWorkgroupSizeX, + kLogicalNotWorkgroupSizeY, + kLogicalNotWorkgroupSizeZ, + }, + { + "logical_or", + kLogicalOrWGSL, + kLogicalOrWorkgroupSizeX, + kLogicalOrWorkgroupSizeY, + kLogicalOrWorkgroupSizeZ, + }, + { + "max_pool2d", + kMaxPool2dWGSL, + kMaxPool2dWorkgroupSizeX, + kMaxPool2dWorkgroupSizeY, + kMaxPool2dWorkgroupSizeZ, + }, + { + "mm", + kMmWGSL, + kMmWorkgroupSizeX, + kMmWorkgroupSizeY, + kMmWorkgroupSizeZ, + }, + { + "mm_tiled", + kMmTiledWGSL, + kMmTiledWorkgroupSizeX, + kMmTiledWorkgroupSizeY, + kMmTiledWorkgroupSizeZ, + }, + { + "mm_vec4", + kMmVec4WGSL, + kMmVec4WorkgroupSizeX, + kMmVec4WorkgroupSizeY, + kMmVec4WorkgroupSizeZ, + }, + { + "native_layer_norm", + kNativeLayerNormWGSL, + kNativeLayerNormWorkgroupSizeX, + kNativeLayerNormWorkgroupSizeY, + kNativeLayerNormWorkgroupSizeZ, + }, + { + "neg", + kNegWGSL, + kNegWorkgroupSizeX, + kNegWorkgroupSizeY, + kNegWorkgroupSizeZ, + }, + { + "permute", + kPermuteWGSL, + kPermuteWorkgroupSizeX, + kPermuteWorkgroupSizeY, + kPermuteWorkgroupSizeZ, + }, + { + "pixel_shuffle", + kPixelShuffleWGSL, + kPixelShuffleWorkgroupSizeX, + kPixelShuffleWorkgroupSizeY, + kPixelShuffleWorkgroupSizeZ, + }, + { + "pow_scalar", + kPowScalarWGSL, + kPowScalarWorkgroupSizeX, + kPowScalarWorkgroupSizeY, + kPowScalarWorkgroupSizeZ, + }, + { + "q4gsw_backward", + kQ4gswBackwardWGSL, + kQ4gswBackwardWorkgroupSizeX, + kQ4gswBackwardWorkgroupSizeY, + kQ4gswBackwardWorkgroupSizeZ, + }, + { + "q4gsw_linear", + kQ4gswLinearWGSL, + kQ4gswLinearWorkgroupSizeX, + kQ4gswLinearWorkgroupSizeY, + kQ4gswLinearWorkgroupSizeZ, + }, + { + "q4gsw_linear_coop4_bicol", + kQ4gswLinearCoop4BicolWGSL, + kQ4gswLinearCoop4BicolWorkgroupSizeX, + kQ4gswLinearCoop4BicolWorkgroupSizeY, + kQ4gswLinearCoop4BicolWorkgroupSizeZ, + }, + { + "q4gsw_linear_gemm_shmem", + kQ4gswLinearGemmShmemWGSL, + kQ4gswLinearGemmShmemWorkgroupSizeX, + kQ4gswLinearGemmShmemWorkgroupSizeY, + kQ4gswLinearGemmShmemWorkgroupSizeZ, + }, + { + "q4gsw_linear_gemm_steel", + kQ4gswLinearGemmSteelWGSL, + kQ4gswLinearGemmSteelWorkgroupSizeX, + kQ4gswLinearGemmSteelWorkgroupSizeY, + kQ4gswLinearGemmSteelWorkgroupSizeZ, + }, + { + "q4gsw_linear_gemm_steel_half", + kQ4gswLinearGemmSteelHalfWGSL, + kQ4gswLinearGemmSteelHalfWorkgroupSizeX, + kQ4gswLinearGemmSteelHalfWorkgroupSizeY, + kQ4gswLinearGemmSteelHalfWorkgroupSizeZ, + }, + { + "q4gsw_linear_gemm_steel_half_pwdq", + kQ4gswLinearGemmSteelHalfPwdqWGSL, + kQ4gswLinearGemmSteelHalfPwdqWorkgroupSizeX, + kQ4gswLinearGemmSteelHalfPwdqWorkgroupSizeY, + kQ4gswLinearGemmSteelHalfPwdqWorkgroupSizeZ, + }, + { + "q4gsw_linear_gemm_steel_half_pwdq_f16acc", + kQ4gswLinearGemmSteelHalfPwdqF16accWGSL, + kQ4gswLinearGemmSteelHalfPwdqF16accWorkgroupSizeX, + kQ4gswLinearGemmSteelHalfPwdqF16accWorkgroupSizeY, + kQ4gswLinearGemmSteelHalfPwdqF16accWorkgroupSizeZ, + }, + { + "q4gsw_qkv_bk64", + kQ4gswQkvBk64WGSL, + kQ4gswQkvBk64WorkgroupSizeX, + kQ4gswQkvBk64WorkgroupSizeY, + kQ4gswQkvBk64WorkgroupSizeZ, + }, + { + "q4gsw_requant", + kQ4gswRequantWGSL, + kQ4gswRequantWorkgroupSizeX, + kQ4gswRequantWorkgroupSizeY, + kQ4gswRequantWorkgroupSizeZ, + }, + { + "q4gsw_steel_bk64", + kQ4gswSteelBk64WGSL, + kQ4gswSteelBk64WorkgroupSizeX, + kQ4gswSteelBk64WorkgroupSizeY, + kQ4gswSteelBk64WorkgroupSizeZ, + }, + { + "q8ta_add", + kQ8taAddWGSL, + kQ8taAddWorkgroupSizeX, + kQ8taAddWorkgroupSizeY, + kQ8taAddWorkgroupSizeZ, + }, + { + "q8ta_conv2d", + kQ8taConv2dWGSL, + kQ8taConv2dWorkgroupSizeX, + kQ8taConv2dWorkgroupSizeY, + kQ8taConv2dWorkgroupSizeZ, + }, + { + "q8ta_conv2d_dw", + kQ8taConv2dDwWGSL, + kQ8taConv2dDwWorkgroupSizeX, + kQ8taConv2dDwWorkgroupSizeY, + kQ8taConv2dDwWorkgroupSizeZ, + }, + { + "q8ta_conv2d_pw", + kQ8taConv2dPwWGSL, + kQ8taConv2dPwWorkgroupSizeX, + kQ8taConv2dPwWorkgroupSizeY, + kQ8taConv2dPwWorkgroupSizeZ, + }, + { + "q8ta_conv2d_transposed", + kQ8taConv2dTransposedWGSL, + kQ8taConv2dTransposedWorkgroupSizeX, + kQ8taConv2dTransposedWorkgroupSizeY, + kQ8taConv2dTransposedWorkgroupSizeZ, + }, + { + "q8ta_linear", + kQ8taLinearWGSL, + kQ8taLinearWorkgroupSizeX, + kQ8taLinearWorkgroupSizeY, + kQ8taLinearWorkgroupSizeZ, + }, + { + "q8ta_pixel_shuffle", + kQ8taPixelShuffleWGSL, + kQ8taPixelShuffleWorkgroupSizeX, + kQ8taPixelShuffleWorkgroupSizeY, + kQ8taPixelShuffleWorkgroupSizeZ, + }, + { + "q8ta_relu", + kQ8taReluWGSL, + kQ8taReluWorkgroupSizeX, + kQ8taReluWorkgroupSizeY, + kQ8taReluWorkgroupSizeZ, + }, + { + "qcs4w_linear", + kQcs4wLinearWGSL, + kQcs4wLinearWorkgroupSizeX, + kQcs4wLinearWorkgroupSizeY, + kQcs4wLinearWorkgroupSizeZ, + }, + { + "quantize_per_tensor", + kQuantizePerTensorWGSL, + kQuantizePerTensorWorkgroupSizeX, + kQuantizePerTensorWorkgroupSizeY, + kQuantizePerTensorWorkgroupSizeZ, + }, + { + "reduce", + kReduceWGSL, + kReduceWorkgroupSizeX, + kReduceWorkgroupSizeY, + kReduceWorkgroupSizeZ, + }, + { + "relu", + kReluWGSL, + kReluWorkgroupSizeX, + kReluWorkgroupSizeY, + kReluWorkgroupSizeZ, + }, + { + "repeat", + kRepeatWGSL, + kRepeatWorkgroupSizeX, + kRepeatWorkgroupSizeY, + kRepeatWorkgroupSizeZ, + }, + { + "rms_norm", + kRmsNormWGSL, + kRmsNormWorkgroupSizeX, + kRmsNormWorkgroupSizeY, + kRmsNormWorkgroupSizeZ, + }, + { + "rms_norm_vec4", + kRmsNormVec4WGSL, + kRmsNormVec4WorkgroupSizeX, + kRmsNormVec4WorkgroupSizeY, + kRmsNormVec4WorkgroupSizeZ, + }, + { + "rotary_embedding", + kRotaryEmbeddingWGSL, + kRotaryEmbeddingWorkgroupSizeX, + kRotaryEmbeddingWorkgroupSizeY, + kRotaryEmbeddingWorkgroupSizeZ, + }, + { + "rotary_embedding_hf", + kRotaryEmbeddingHfWGSL, + kRotaryEmbeddingHfWorkgroupSizeX, + kRotaryEmbeddingHfWorkgroupSizeY, + kRotaryEmbeddingHfWorkgroupSizeZ, + }, + { + "round", + kRoundWGSL, + kRoundWorkgroupSizeX, + kRoundWorkgroupSizeY, + kRoundWorkgroupSizeZ, + }, + { + "rsqrt", + kRsqrtWGSL, + kRsqrtWorkgroupSizeX, + kRsqrtWorkgroupSizeY, + kRsqrtWorkgroupSizeZ, + }, + { + "sdpa_compute_attn_weights", + kSdpaComputeAttnWeightsWGSL, + kSdpaComputeAttnWeightsWorkgroupSizeX, + kSdpaComputeAttnWeightsWorkgroupSizeY, + kSdpaComputeAttnWeightsWorkgroupSizeZ, + }, + { + "sdpa_compute_attn_weights_half", + kSdpaComputeAttnWeightsHalfWGSL, + kSdpaComputeAttnWeightsHalfWorkgroupSizeX, + kSdpaComputeAttnWeightsHalfWorkgroupSizeY, + kSdpaComputeAttnWeightsHalfWorkgroupSizeZ, + }, + { + "sdpa_compute_out", + kSdpaComputeOutWGSL, + kSdpaComputeOutWorkgroupSizeX, + kSdpaComputeOutWorkgroupSizeY, + kSdpaComputeOutWorkgroupSizeZ, + }, + { + "sdpa_compute_out_half", + kSdpaComputeOutHalfWGSL, + kSdpaComputeOutHalfWorkgroupSizeX, + kSdpaComputeOutHalfWorkgroupSizeY, + kSdpaComputeOutHalfWorkgroupSizeZ, + }, + { + "sdpa_fd_reduce", + kSdpaFdReduceWGSL, + kSdpaFdReduceWorkgroupSizeX, + kSdpaFdReduceWorkgroupSizeY, + kSdpaFdReduceWorkgroupSizeZ, + }, + { + "sdpa_fd_split", + kSdpaFdSplitWGSL, + kSdpaFdSplitWorkgroupSizeX, + kSdpaFdSplitWorkgroupSizeY, + kSdpaFdSplitWorkgroupSizeZ, + }, + { + "sdpa_fd_split_half", + kSdpaFdSplitHalfWGSL, + kSdpaFdSplitHalfWorkgroupSizeX, + kSdpaFdSplitHalfWorkgroupSizeY, + kSdpaFdSplitHalfWorkgroupSizeZ, + }, + { + "sdpa_softmax", + kSdpaSoftmaxWGSL, + kSdpaSoftmaxWorkgroupSizeX, + kSdpaSoftmaxWorkgroupSizeY, + kSdpaSoftmaxWorkgroupSizeZ, + }, + { + "select", + kSelectWGSL, + kSelectWorkgroupSizeX, + kSelectWorkgroupSizeY, + kSelectWorkgroupSizeZ, + }, + { + "sigmoid", + kSigmoidWGSL, + kSigmoidWorkgroupSizeX, + kSigmoidWorkgroupSizeY, + kSigmoidWorkgroupSizeZ, + }, + { + "silu_mul_fused", + kSiluMulFusedWGSL, + kSiluMulFusedWorkgroupSizeX, + kSiluMulFusedWorkgroupSizeY, + kSiluMulFusedWorkgroupSizeZ, + }, + { + "sin", + kSinWGSL, + kSinWorkgroupSizeX, + kSinWorkgroupSizeY, + kSinWorkgroupSizeZ, + }, + { + "slice", + kSliceWGSL, + kSliceWorkgroupSizeX, + kSliceWorkgroupSizeY, + kSliceWorkgroupSizeZ, + }, + { + "softmax", + kSoftmaxWGSL, + kSoftmaxWorkgroupSizeX, + kSoftmaxWorkgroupSizeY, + kSoftmaxWorkgroupSizeZ, + }, + { + "sqrt", + kSqrtWGSL, + kSqrtWorkgroupSizeX, + kSqrtWorkgroupSizeY, + kSqrtWorkgroupSizeZ, + }, + { + "streaming_attention_k16_causal_bound", + kStreamingAttentionK16CausalBoundWGSL, + kStreamingAttentionK16CausalBoundWorkgroupSizeX, + kStreamingAttentionK16CausalBoundWorkgroupSizeY, + kStreamingAttentionK16CausalBoundWorkgroupSizeZ, + }, + { + "streaming_attention_qwen3_k16_causal_bound", + kStreamingAttentionQwen3K16CausalBoundWGSL, + kStreamingAttentionQwen3K16CausalBoundWorkgroupSizeX, + kStreamingAttentionQwen3K16CausalBoundWorkgroupSizeY, + kStreamingAttentionQwen3K16CausalBoundWorkgroupSizeZ, + }, + { + "streaming_attention_qwen3_q32_k16_causal_bound", + kStreamingAttentionQwen3Q32K16CausalBoundWGSL, + kStreamingAttentionQwen3Q32K16CausalBoundWorkgroupSizeX, + kStreamingAttentionQwen3Q32K16CausalBoundWorkgroupSizeY, + kStreamingAttentionQwen3Q32K16CausalBoundWorkgroupSizeZ, + }, + { + "tanh", + kTanhWGSL, + kTanhWorkgroupSizeX, + kTanhWorkgroupSizeY, + kTanhWorkgroupSizeZ, + }, + { + "to_copy_bool_to_float", + kToCopyBoolToFloatWGSL, + kToCopyBoolToFloatWorkgroupSizeX, + kToCopyBoolToFloatWorkgroupSizeY, + kToCopyBoolToFloatWorkgroupSizeZ, + }, + { + "to_copy_float_to_int", + kToCopyFloatToIntWGSL, + kToCopyFloatToIntWorkgroupSizeX, + kToCopyFloatToIntWorkgroupSizeY, + kToCopyFloatToIntWorkgroupSizeZ, + }, + { + "to_copy_int_to_float", + kToCopyIntToFloatWGSL, + kToCopyIntToFloatWorkgroupSizeX, + kToCopyIntToFloatWorkgroupSizeY, + kToCopyIntToFloatWorkgroupSizeZ, + }, + { + "update_cache", + kUpdateCacheWGSL, + kUpdateCacheWorkgroupSizeX, + kUpdateCacheWorkgroupSizeY, + kUpdateCacheWorkgroupSizeZ, + }, + { + "update_cache_half", + kUpdateCacheHalfWGSL, + kUpdateCacheHalfWorkgroupSizeX, + kUpdateCacheHalfWorkgroupSizeY, + kUpdateCacheHalfWorkgroupSizeZ, + }, + { + "upsample_bilinear2d", + kUpsampleBilinear2dWGSL, + kUpsampleBilinear2dWorkgroupSizeX, + kUpsampleBilinear2dWorkgroupSizeY, + kUpsampleBilinear2dWorkgroupSizeZ, + }, + { + "upsample_nearest2d", + kUpsampleNearest2dWGSL, + kUpsampleNearest2dWorkgroupSizeX, + kUpsampleNearest2dWorkgroupSizeY, + kUpsampleNearest2dWorkgroupSizeZ, + }, + { + "where", + kWhereWGSL, + kWhereWorkgroupSizeX, + kWhereWorkgroupSizeY, + kWhereWorkgroupSizeZ, + }, +}}; + +} // namespace + +const WebGPUShaderInfo& get_webgpu_shader_info(std::string_view name) { + for (const auto& shader : kShaderRegistry) { + if (shader.name == name) { + return shader; + } + } + throw std::runtime_error( + "WebGPU shader registry: unknown shader '" + std::string(name) + "'"); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/WebGPUShaderRegistry.h b/backends/webgpu/runtime/WebGPUShaderRegistry.h new file mode 100644 index 00000000000..64b86e36110 --- /dev/null +++ b/backends/webgpu/runtime/WebGPUShaderRegistry.h @@ -0,0 +1,26 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include +#include + +namespace executorch::backends::webgpu { + +struct WebGPUShaderInfo { + std::string_view name; + const char* source; + uint32_t workgroup_size_x; + uint32_t workgroup_size_y; + uint32_t workgroup_size_z; +}; + +const WebGPUShaderInfo& get_webgpu_shader_info(std::string_view name); + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/WebGPUUtils.h b/backends/webgpu/runtime/WebGPUUtils.h index 06c7f312dcd..20ff14c34f1 100644 --- a/backends/webgpu/runtime/WebGPUUtils.h +++ b/backends/webgpu/runtime/WebGPUUtils.h @@ -16,8 +16,10 @@ #include #include #include +#include #include #include +#include #include #include #include @@ -30,6 +32,18 @@ inline uint64_t numel_of(const std::vector& dims) { return numel(dims); } +// fp32, non-null-buffer tensor with byte size matching its element count; +// the dtype/aliasing precondition fusion passes require of every operand. +inline bool is_fp32_tensor(const WebGPUTensor& tensor) { + if (tensor.is_int || tensor.elem_size != sizeof(float) || + tensor.buffer == nullptr) { + return false; + } + const uint64_t elems = numel_of(tensor.dims); + return elems <= std::numeric_limits::max() / sizeof(float) && + tensor.nbytes == static_cast(elems) * sizeof(float); +} + // Clamp workgroup size to device limit (SwiftShader caps at 128). inline uint32_t clamp_workgroup_size(WGPUDevice device, uint32_t desired) { WGPULimits limits = {}; @@ -50,11 +64,6 @@ inline uint32_t clamp_workgroup_size_pow2(WGPUDevice device, uint32_t desired) { return p; } -struct WgCount { - uint32_t x; - uint32_t y; -}; - // Device's max workgroups per dispatch dimension; the WebGPU spec-default floor // (65535) if the query fails — never under-reports a real device's capacity. inline uint32_t queried_max_workgroups(WGPUDevice device) { @@ -514,6 +523,53 @@ inline ComputePipelineBundle make_compute_pipeline( return bundle; } +// Builds a pipeline for a different shader that uses the exact layout and +// bind group of an earlier bundle. Multi-route ops use this when alternate +// shaders have an identical binding contract: only the shader module and +// pipeline are new; layout and bind-group construction stay single-copy. +inline ComputePipelineBundle make_compute_pipeline( + WGPUDevice device, + const char* wgsl_source, + const ComputePipelineBundle& shared_resources, + const WGPUConstantEntry* constants = nullptr, + size_t constant_count = 0, + const char* entry_point = "main") { + if (shared_resources.bind_group_layout == nullptr || + shared_resources.pipeline_layout == nullptr || + shared_resources.bind_group == nullptr) { + throw std::runtime_error( + "make_compute_pipeline: shared resources are not available"); + } + + ComputePipelineBundle bundle; + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {wgsl_source, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + bundle.shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + if (bundle.shader == nullptr) { + throw std::runtime_error( + "make_compute_pipeline: shader module creation failed"); + } + + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.layout = shared_resources.pipeline_layout; + pipeline_desc.compute.module = bundle.shader; + pipeline_desc.compute.entryPoint = {entry_point, WGPU_STRLEN}; + pipeline_desc.compute.constantCount = constant_count; + pipeline_desc.compute.constants = constants; + bundle.pipeline = wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + if (bundle.pipeline == nullptr) { + throw std::runtime_error( + "make_compute_pipeline: compute pipeline creation failed"); + } + + wgpuBindGroupAddRef(shared_resources.bind_group); + bundle.bind_group = shared_resources.bind_group; + return bundle; +} + // The {wg_size, stride_x} override-constant pair every 2D-spill dispatch // builds from its DispatchGrid; was hand-rolled identically at 7 call sites. inline std::array make_grid_constants( diff --git a/backends/webgpu/runtime/ops/amax/Reduce.cpp b/backends/webgpu/runtime/ops/amax/Reduce.cpp index 14675d97e96..c1dff62e035 100644 --- a/backends/webgpu/runtime/ops/amax/Reduce.cpp +++ b/backends/webgpu/runtime/ops/amax/Reduce.cpp @@ -9,7 +9,7 @@ #include #include #include -#include +#include #include diff --git a/backends/webgpu/runtime/ops/amin/Reduce.cpp b/backends/webgpu/runtime/ops/amin/Reduce.cpp index fbe574fdf0b..24a5ebeb826 100644 --- a/backends/webgpu/runtime/ops/amin/Reduce.cpp +++ b/backends/webgpu/runtime/ops/amin/Reduce.cpp @@ -9,7 +9,7 @@ #include #include #include -#include +#include #include diff --git a/backends/webgpu/runtime/ops/amin/amin.wgsl b/backends/webgpu/runtime/ops/amin/amin.wgsl deleted file mode 100644 index 4778800ab3d..00000000000 --- a/backends/webgpu/runtime/ops/amin/amin.wgsl +++ /dev/null @@ -1,49 +0,0 @@ -@group(0) @binding(0) var input: array; -@group(0) @binding(1) var output: array; - -struct Params { - num_rows: u32, - reduce_size: u32, -} -@group(0) @binding(2) var params: Params; - -override wg_size: u32 = 256u; - -// Cooperative shared-memory reduction; mirrors Vulkan reduce.glsl (a group of -// threads co-operates per reduction row, partials aggregated in shared memory). -// Fixed upper bound (>= any clamped wg_size); only [0, wg_size) is used. -var partials: array; - -@compute @workgroup_size(wg_size) -fn main( - @builtin(workgroup_id) wid: vec3, - @builtin(local_invocation_id) lid: vec3, - @builtin(num_workgroups) num_workgroups: vec3) { - // One workgroup per reduction row; 2D-fold lifts the 65535 grid cap. - let row = wid.x + wid.y * num_workgroups.x; - if (row >= params.num_rows) { - return; - } - let base = row * params.reduce_size; - - // Each thread reduces a strided slice of the row into a partial. Seed with - // the row's first element (always valid; reduce_size >= 1) so threads that - // own no element contribute a real value, not an out-of-range identity. - var acc = input[base]; - var i = lid.x; - while (i < params.reduce_size) { - acc = min(acc, input[base + i]); - i = i + wg_size; - } - partials[lid.x] = acc; - workgroupBarrier(); - - // Thread 0 aggregates the wg_size partials (mirrors Vulkan's group aggregate). - if (lid.x == 0u) { - var m = partials[0]; - for (var t = 1u; t < wg_size; t = t + 1u) { - m = min(m, partials[t]); - } - output[row] = m; - } -} diff --git a/backends/webgpu/runtime/ops/floor_divide/binary_floor_divide_wgsl.h b/backends/webgpu/runtime/ops/binary_op/binary_floor_divide_wgsl.h similarity index 97% rename from backends/webgpu/runtime/ops/floor_divide/binary_floor_divide_wgsl.h rename to backends/webgpu/runtime/ops/binary_op/binary_floor_divide_wgsl.h index fe2315df30f..dcd2b2e46d8 100644 --- a/backends/webgpu/runtime/ops/floor_divide/binary_floor_divide_wgsl.h +++ b/backends/webgpu/runtime/ops/binary_op/binary_floor_divide_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from binary_floor_divide.wgsl - DO NOT EDIT. +// @generated from binary_op.wgsl - DO NOT EDIT. // wgsl-sha256: baf71d277da79389315a6b96b439e7f0a55842e8288283f2af121f84536b3af3 inline constexpr const char* kBinaryFloorDivideWGSL = R"( @group(0) @binding(0) var input1: array; diff --git a/backends/webgpu/runtime/ops/minimum/binary_minimum_wgsl.h b/backends/webgpu/runtime/ops/binary_op/binary_minimum_wgsl.h similarity index 97% rename from backends/webgpu/runtime/ops/minimum/binary_minimum_wgsl.h rename to backends/webgpu/runtime/ops/binary_op/binary_minimum_wgsl.h index 88d9614ba59..c3f5f76c81f 100644 --- a/backends/webgpu/runtime/ops/minimum/binary_minimum_wgsl.h +++ b/backends/webgpu/runtime/ops/binary_op/binary_minimum_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from binary_minimum.wgsl - DO NOT EDIT. +// @generated from binary_op.wgsl - DO NOT EDIT. // wgsl-sha256: 929b7ba85936e3652baea9f4e5e7f049d232c7ae7a74814a536b4c2674897972 inline constexpr const char* kBinaryMinimumWGSL = R"( @group(0) @binding(0) var input1: array; diff --git a/backends/webgpu/runtime/ops/mul/binary_mul_wgsl.h b/backends/webgpu/runtime/ops/binary_op/binary_mul_wgsl.h similarity index 98% rename from backends/webgpu/runtime/ops/mul/binary_mul_wgsl.h rename to backends/webgpu/runtime/ops/binary_op/binary_mul_wgsl.h index c9f60dbd200..68784e82eb2 100644 --- a/backends/webgpu/runtime/ops/mul/binary_mul_wgsl.h +++ b/backends/webgpu/runtime/ops/binary_op/binary_mul_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from binary_mul.wgsl - DO NOT EDIT. +// @generated from binary_op.wgsl - DO NOT EDIT. // wgsl-sha256: d248c0f1856b57115a5001a47f4936caa564dd3b787c02ceba504a13ab987812 inline constexpr const char* kBinaryMulWGSL = R"( @group(0) @binding(0) var input1: array; diff --git a/backends/webgpu/runtime/ops/binary_op/binary_op.wgsl b/backends/webgpu/runtime/ops/binary_op/binary_op.wgsl index 98076ed98b5..4b42665013f 100644 --- a/backends/webgpu/runtime/ops/binary_op/binary_op.wgsl +++ b/backends/webgpu/runtime/ops/binary_op/binary_op.wgsl @@ -16,11 +16,14 @@ override wg_size: u32 = 64u; $if USE_ALPHA: override alpha: f32 = 1.0; -fn op(a: f32, b: f32) -> f32 { - return ${OP_EXPR}; -} +$if INLINE: + @compute @workgroup_size(wg_size, 1, 1) +$else: + fn op(a: f32, b: f32) -> f32 { + return ${OP_EXPR}; + } -@compute @workgroup_size(wg_size, 1, 1) + @compute @workgroup_size(wg_size, 1, 1) fn main( @builtin(global_invocation_id) gid: vec3, @builtin(num_workgroups) num_workgroups: vec3) { @@ -30,6 +33,8 @@ fn main( return; } + $if INLINE: + // Fast path: every input dim matches the output dim -> elementwise. var same = true; for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { if (in1_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u] || @@ -38,10 +43,15 @@ fn main( } } if (same) { - output[idx] = op(input1[idx], input2[idx]); + $if INLINE: + output[idx] = ${SAME_EXPR}; + $else: + output[idx] = op(input1[idx], input2[idx]); return; } + $if INLINE: + // Broadcast: out idx -> per-input coord (clamp size-1 dims), relinearize. var rem = idx; var l1: u32 = 0u; var l2: u32 = 0u; @@ -51,5 +61,8 @@ fn main( l1 = l1 + min(coord, in1_meta.sizes[d >> 2u][d & 3u] - 1u) * in1_meta.strides[d >> 2u][d & 3u]; l2 = l2 + min(coord, in2_meta.sizes[d >> 2u][d & 3u] - 1u) * in2_meta.strides[d >> 2u][d & 3u]; } - output[idx] = op(input1[l1], input2[l2]); + $if INLINE: + output[idx] = ${BROADCAST_EXPR}; + $else: + output[idx] = op(input1[l1], input2[l2]); } diff --git a/backends/webgpu/runtime/ops/binary_op/binary_op.yaml b/backends/webgpu/runtime/ops/binary_op/binary_op.yaml index 03040fcbcce..bd4cd452d35 100644 --- a/backends/webgpu/runtime/ops/binary_op/binary_op.yaml +++ b/backends/webgpu/runtime/ops/binary_op/binary_op.yaml @@ -2,6 +2,9 @@ binary_op: parameter_names_with_default_values: OP_EXPR: a + alpha * b USE_ALPHA: 1 + INLINE: 0 + SAME_EXPR: input1[idx] + input2[idx] + BROADCAST_EXPR: input1[l1] + input2[l2] shader_variants: - NAME: binary_div OP_EXPR: a / b @@ -9,3 +12,23 @@ binary_op: - NAME: binary_sub OP_EXPR: a - alpha * b USE_ALPHA: 1 + - NAME: binary_minimum + USE_ALPHA: 0 + INLINE: 1 + SAME_EXPR: min(input1[idx], input2[idx]) + BROADCAST_EXPR: min(input1[l1], input2[l2]) + - NAME: binary_pow + USE_ALPHA: 0 + INLINE: 1 + SAME_EXPR: pow(input1[idx], input2[idx]) + BROADCAST_EXPR: pow(input1[l1], input2[l2]) + - NAME: binary_floor_divide + USE_ALPHA: 0 + INLINE: 1 + SAME_EXPR: floor(input1[idx] / input2[idx]) + BROADCAST_EXPR: floor(input1[l1] / input2[l2]) + - NAME: binary_mul + USE_ALPHA: 0 + INLINE: 1 + SAME_EXPR: input1[idx] * input2[idx] + BROADCAST_EXPR: input1[l1] * input2[l2] diff --git a/backends/webgpu/runtime/ops/pow/binary_pow_wgsl.h b/backends/webgpu/runtime/ops/binary_op/binary_pow_wgsl.h similarity index 98% rename from backends/webgpu/runtime/ops/pow/binary_pow_wgsl.h rename to backends/webgpu/runtime/ops/binary_op/binary_pow_wgsl.h index 3532c091160..776d4b6693f 100644 --- a/backends/webgpu/runtime/ops/pow/binary_pow_wgsl.h +++ b/backends/webgpu/runtime/ops/binary_op/binary_pow_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from binary_pow.wgsl - DO NOT EDIT. +// @generated from binary_op.wgsl - DO NOT EDIT. // wgsl-sha256: a88c161bd3f43d21a72ebd8ca6f8611b6b9b854e3572a8e6b820602091bc464c inline constexpr const char* kBinaryPowWGSL = R"( @group(0) @binding(0) var input1: array; diff --git a/backends/webgpu/runtime/ops/compare/Compare.cpp b/backends/webgpu/runtime/ops/compare/Compare.cpp index 99200cb2425..f47f9491191 100644 --- a/backends/webgpu/runtime/ops/compare/Compare.cpp +++ b/backends/webgpu/runtime/ops/compare/Compare.cpp @@ -58,13 +58,13 @@ void compare_impl( in2_tensor.elem_size != 4) { throw std::runtime_error("compare: fp32 inputs only"); } - if (!out_tensor.is_int || out_tensor.elem_size != 1) { + if (!out_tensor.is_bool || out_tensor.elem_size != 1) { throw std::runtime_error("compare: output must be a 1-byte bool tensor"); } const uint64_t numel = out_tensor.nbytes; - // out bool packed 4/word (array); numel%4==0 gates the readback map. - if (numel == 0u || numel % 4u != 0u || numel > UINT32_MAX) { - throw std::runtime_error("compare: numel must be a nonzero mult of 4"); + // out bool is byte-packed into ceil(numel / 4) u32 storage words. + if (numel == 0u || numel > UINT32_MAX) { + throw std::runtime_error("compare: numel must be nonzero and fit u32"); } const uint64_t in_numel = in1_tensor.nbytes / sizeof(float); if (in1_tensor.nbytes != in2_tensor.nbytes || in_numel != numel) { @@ -75,7 +75,7 @@ void compare_impl( params.num_elements = static_cast(numel); params.op = op; - const uint32_t words = static_cast(numel / 4u); + const uint32_t words = static_cast((numel + 3u) / 4u); uint32_t wg_size = utils::clamp_workgroup_size(device, kCompareWorkgroupSizeX); utils::WgCount workgroup_count = @@ -97,7 +97,7 @@ void compare_impl( {0, WGPUBufferBindingType_Storage, out_tensor.buffer, - out_tensor.nbytes}, + static_cast(words) * sizeof(uint32_t)}, {1, WGPUBufferBindingType_ReadOnlyStorage, in1_tensor.buffer, @@ -127,9 +127,8 @@ void compare_impl( WebGPUGraph& g) { const auto& d = g.cur_dims(in1_id); const uint64_t n = utils::numel_of(d); - if (n == 0u || n % 4u != 0u || n > UINT32_MAX || - utils::numel_of(g.cur_dims(in2_id)) != n) { - throw std::runtime_error("compare(resize): numel must be a mult of 4"); + if (n == 0u || n > UINT32_MAX || utils::numel_of(g.cur_dims(in2_id)) != n) { + throw std::runtime_error("compare(resize): invalid numel"); } g.set_cur_dims(out_id, d); CompareParams p = {}; @@ -137,7 +136,7 @@ void compare_impl( p.op = op; wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); const utils::WgCount wgc = utils::compute_2d_workgroup_count( - g.device(), static_cast(n / 4u), wg_size, "compare"); + g.device(), static_cast((n + 3u) / 4u), wg_size, "compare"); g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x; g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y; }; diff --git a/backends/webgpu/runtime/ops/compare/compare.wgsl b/backends/webgpu/runtime/ops/compare/compare.wgsl index a16568e5533..a650a6298b5 100644 --- a/backends/webgpu/runtime/ops/compare/compare.wgsl +++ b/backends/webgpu/runtime/ops/compare/compare.wgsl @@ -16,27 +16,29 @@ override wg_size: u32 = 64u; fn main( @builtin(global_invocation_id) gid: vec3, @builtin(num_workgroups) num_workgroups: vec3) { - // One thread per output word = 4 bool bytes; num_elements%4==0 (host). + // One thread per output word = up to 4 bool bytes. let widx = gid.x + gid.y * (num_workgroups.x * wg_size); - let words = (params.num_elements + 3u) / 4u; + let words = (params.num_elements - 1u) / 4u + 1u; if (widx >= words) { return; } var packed: u32 = 0u; for (var j: u32 = 0u; j < 4u; j = j + 1u) { let i = widx * 4u + j; - let a = input1[i]; - let b = input2[i]; - var r: bool; - switch params.op { - case 0u: { r = a == b; } // eq - case 1u: { r = a < b; } // lt - case 2u: { r = a <= b; } // le - case 3u: { r = a > b; } // gt - default: { r = a >= b; } // ge - } - if (r) { - packed = packed | (1u << (j * 8u)); + if (i < params.num_elements) { + let a = input1[i]; + let b = input2[i]; + var r: bool; + switch params.op { + case 0u: { r = a == b; } // eq + case 1u: { r = a < b; } // lt + case 2u: { r = a <= b; } // le + case 3u: { r = a > b; } // gt + default: { r = a >= b; } // ge + } + if (r) { + packed = packed | (1u << (j * 8u)); + } } } t_out[widx] = packed; diff --git a/backends/webgpu/runtime/ops/compare/compare_wgsl.h b/backends/webgpu/runtime/ops/compare/compare_wgsl.h index 672c99b62d8..c1c4ac23e4f 100644 --- a/backends/webgpu/runtime/ops/compare/compare_wgsl.h +++ b/backends/webgpu/runtime/ops/compare/compare_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from compare.wgsl - DO NOT EDIT. -// wgsl-sha256: 241e7e6762b1eded07d28a3767936c970f509f6591e7cbf599d0b1eb61efb181 +// wgsl-sha256: 8f330b5a1e29a1fb8135e64600dadb1dc64c98f8f5371357f8de73a1808b76d6 inline constexpr const char* kCompareWGSL = R"( @group(0) @binding(0) var t_out: array; @group(0) @binding(1) var input1: array; @@ -33,27 +33,29 @@ override wg_size: u32 = 64u; fn main( @builtin(global_invocation_id) gid: vec3, @builtin(num_workgroups) num_workgroups: vec3) { - // One thread per output word = 4 bool bytes; num_elements%4==0 (host). + // One thread per output word = up to 4 bool bytes. let widx = gid.x + gid.y * (num_workgroups.x * wg_size); - let words = (params.num_elements + 3u) / 4u; + let words = (params.num_elements - 1u) / 4u + 1u; if (widx >= words) { return; } var packed: u32 = 0u; for (var j: u32 = 0u; j < 4u; j = j + 1u) { let i = widx * 4u + j; - let a = input1[i]; - let b = input2[i]; - var r: bool; - switch params.op { - case 0u: { r = a == b; } // eq - case 1u: { r = a < b; } // lt - case 2u: { r = a <= b; } // le - case 3u: { r = a > b; } // gt - default: { r = a >= b; } // ge - } - if (r) { - packed = packed | (1u << (j * 8u)); + if (i < params.num_elements) { + let a = input1[i]; + let b = input2[i]; + var r: bool; + switch params.op { + case 0u: { r = a == b; } // eq + case 1u: { r = a < b; } // lt + case 2u: { r = a <= b; } // le + case 3u: { r = a > b; } // gt + default: { r = a >= b; } // ge + } + if (r) { + packed = packed | (1u << (j * 8u)); + } } } t_out[widx] = packed; diff --git a/backends/webgpu/runtime/ops/embedding_q4gsw/EmbeddingQ4gsw.cpp b/backends/webgpu/runtime/ops/embedding_q4gsw/EmbeddingQ4gsw.cpp index 64956121b17..b5ea0cdca4c 100644 --- a/backends/webgpu/runtime/ops/embedding_q4gsw/EmbeddingQ4gsw.cpp +++ b/backends/webgpu/runtime/ops/embedding_q4gsw/EmbeddingQ4gsw.cpp @@ -37,13 +37,37 @@ static_assert( sizeof(EmbeddingParams) == 32, "EmbeddingParams must be 32 bytes"); +struct EmbeddingLayout { + uint32_t embed_dim; + uint32_t blocks_per_row; + uint32_t group_size; + uint32_t groups_per_row; + uint32_t bytes_per_row; + bool is_linear_weight; +}; + +EmbeddingParams make_embedding_params( + const EmbeddingLayout& layout, + uint32_t num_indices, + uint32_t total_blocks) { + return { + layout.embed_dim, + layout.blocks_per_row, + num_indices, + layout.group_size, + layout.groups_per_row, + layout.bytes_per_row, + total_blocks, + layout.is_linear_weight ? 1u : 0u}; +} + // Resize hook body: recompute counts/dispatch; out = indices dims + // [embed_dim]. void resize_embedding_q4gsw( WebGPUGraph& g, int indices_id, int out_id, - EmbeddingParams params, + const EmbeddingLayout& layout, uint32_t wg_size, size_t dispatch_idx, WGPUBuffer params_buf) { @@ -52,17 +76,17 @@ void resize_embedding_q4gsw( if (ni == 0) { throw std::runtime_error("WebGPU embedding_q4gsw: zero indices"); } - const uint64_t total_blocks = ni * params.blocks_per_row; + const uint64_t total_blocks = ni * layout.blocks_per_row; if (total_blocks > UINT32_MAX) { throw std::runtime_error( "WebGPU embedding_q4gsw: total_blocks exceeds uint32"); } std::vector od = id; - od.push_back(static_cast(params.embed_dim)); + od.push_back(static_cast(layout.embed_dim)); g.set_cur_dims(out_id, od); - params.num_indices = static_cast(ni); - params.total_blocks = static_cast(total_blocks); - wgpuQueueWriteBuffer(g.queue(), params_buf, 0, ¶ms, sizeof(params)); + EmbeddingParams p = make_embedding_params( + layout, static_cast(ni), static_cast(total_blocks)); + wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); g.dispatch_at(dispatch_idx).workgroup_count_x = utils::compute_1d_workgroup_count( g.device(), @@ -166,15 +190,15 @@ void embedding_q4gsw_impl(WebGPUGraph& graph, const std::vector& args) { const uint32_t workgroup_count = utils::compute_1d_workgroup_count( device, static_cast(total_blocks), wg_size, "embedding_q4gsw"); - EmbeddingParams params = {}; - params.embed_dim = embed_dim; - params.blocks_per_row = blocks_per_row; - params.num_indices = num_indices; // std140 layout only; shader derives it - params.group_size = static_cast(group_size); - params.groups_per_row = groups_per_row; - params.bytes_per_row = bytes_per_row; - params.total_blocks = static_cast(total_blocks); - params.is_linear_weight = is_linear ? 1u : 0u; + const EmbeddingLayout layout = { + embed_dim, + blocks_per_row, + static_cast(group_size), + groups_per_row, + bytes_per_row, + is_linear}; + EmbeddingParams params = make_embedding_params( + layout, num_indices, static_cast(total_blocks)); WGPUBufferDescriptor uniform_desc = {}; uniform_desc.size = sizeof(EmbeddingParams); @@ -223,10 +247,10 @@ void embedding_q4gsw_impl(WebGPUGraph& graph, const std::vector& args) { WGPUBuffer params_buf = uniform_buffer; graph.add_tensor_resize_hook( indices_id, - [indices_id, out_id, params, wg_size, dispatch_idx, params_buf]( + [indices_id, out_id, layout, wg_size, dispatch_idx, params_buf]( WebGPUGraph& g) { resize_embedding_q4gsw( - g, indices_id, out_id, params, wg_size, dispatch_idx, params_buf); + g, indices_id, out_id, layout, wg_size, dispatch_idx, params_buf); }); // Graph owns it so the resize hook can rewrite it; freed in the dtor. diff --git a/backends/webgpu/runtime/ops/amax/amax_wgsl.h b/backends/webgpu/runtime/ops/extrema/amax_wgsl.h similarity index 97% rename from backends/webgpu/runtime/ops/amax/amax_wgsl.h rename to backends/webgpu/runtime/ops/extrema/amax_wgsl.h index 48ec8f20e27..0b7a4d2238e 100644 --- a/backends/webgpu/runtime/ops/amax/amax_wgsl.h +++ b/backends/webgpu/runtime/ops/extrema/amax_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from amax.wgsl - DO NOT EDIT. +// @generated from extrema.wgsl - DO NOT EDIT. // wgsl-sha256: 35fc059d7c72caa17f9cb1128823ecfd8f75be4ce24b6cd4f9629a97b52f64c0 inline constexpr const char* kAmaxWGSL = R"( @group(0) @binding(0) var input: array; diff --git a/backends/webgpu/runtime/ops/amin/amin_wgsl.h b/backends/webgpu/runtime/ops/extrema/amin_wgsl.h similarity index 97% rename from backends/webgpu/runtime/ops/amin/amin_wgsl.h rename to backends/webgpu/runtime/ops/extrema/amin_wgsl.h index 40a97c67a63..8b8bf5456b1 100644 --- a/backends/webgpu/runtime/ops/amin/amin_wgsl.h +++ b/backends/webgpu/runtime/ops/extrema/amin_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from amin.wgsl - DO NOT EDIT. +// @generated from extrema.wgsl - DO NOT EDIT. // wgsl-sha256: 8cb6035ae4d34eb2a6cc973d93d9847905722e967239c96033fccfe3a1943cb2 inline constexpr const char* kAminWGSL = R"( @group(0) @binding(0) var input: array; diff --git a/backends/webgpu/runtime/ops/amax/amax.wgsl b/backends/webgpu/runtime/ops/extrema/extrema.wgsl similarity index 94% rename from backends/webgpu/runtime/ops/amax/amax.wgsl rename to backends/webgpu/runtime/ops/extrema/extrema.wgsl index 2f23b9c38fa..6c55d60d9b7 100644 --- a/backends/webgpu/runtime/ops/amax/amax.wgsl +++ b/backends/webgpu/runtime/ops/extrema/extrema.wgsl @@ -32,7 +32,7 @@ fn main( var acc = input[base]; var i = lid.x; while (i < params.reduce_size) { - acc = max(acc, input[base + i]); + acc = ${REDUCE_FN}(acc, input[base + i]); i = i + wg_size; } partials[lid.x] = acc; @@ -42,7 +42,7 @@ fn main( if (lid.x == 0u) { var m = partials[0]; for (var t = 1u; t < wg_size; t = t + 1u) { - m = max(m, partials[t]); + m = ${REDUCE_FN}(m, partials[t]); } output[row] = m; } diff --git a/backends/webgpu/runtime/ops/extrema/extrema.yaml b/backends/webgpu/runtime/ops/extrema/extrema.yaml new file mode 100644 index 00000000000..a85da849c60 --- /dev/null +++ b/backends/webgpu/runtime/ops/extrema/extrema.yaml @@ -0,0 +1,13 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +extrema: + parameter_names_with_default_values: + REDUCE_FN: max + shader_variants: + - NAME: amax + - NAME: amin + REDUCE_FN: min diff --git a/backends/webgpu/runtime/ops/floor_divide/BinaryOp.cpp b/backends/webgpu/runtime/ops/floor_divide/BinaryOp.cpp index a35a7587311..98a76228003 100644 --- a/backends/webgpu/runtime/ops/floor_divide/BinaryOp.cpp +++ b/backends/webgpu/runtime/ops/floor_divide/BinaryOp.cpp @@ -9,7 +9,7 @@ #include #include #include -#include +#include #include #include diff --git a/backends/webgpu/runtime/ops/floor_divide/binary_floor_divide.wgsl b/backends/webgpu/runtime/ops/floor_divide/binary_floor_divide.wgsl deleted file mode 100644 index 3b4edd788b0..00000000000 --- a/backends/webgpu/runtime/ops/floor_divide/binary_floor_divide.wgsl +++ /dev/null @@ -1,51 +0,0 @@ -@group(0) @binding(0) var input1: array; -@group(0) @binding(1) var input2: array; -@group(0) @binding(2) var output: array; - -struct TensorMeta { - ndim: u32, - numel: u32, - sizes: array, 2>, - strides: array, 2>, -} -@group(0) @binding(3) var out_meta: TensorMeta; -@group(0) @binding(4) var in1_meta: TensorMeta; -@group(0) @binding(5) var in2_meta: TensorMeta; - -override wg_size: u32 = 64u; - -@compute @workgroup_size(wg_size, 1, 1) -fn main( - @builtin(global_invocation_id) gid: vec3, - @builtin(num_workgroups) num_workgroups: vec3) { - // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel). - let idx = gid.x + gid.y * (num_workgroups.x * wg_size); - if (idx >= out_meta.numel) { - return; - } - - // Fast path: every input dim matches the output dim -> elementwise. - var same = true; - for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - if (in1_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u] || - in2_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u]) { - same = false; - } - } - if (same) { - output[idx] = floor(input1[idx] / input2[idx]); - return; - } - - // Broadcast: out idx -> per-input coord (clamp size-1 dims), relinearize. - var rem = idx; - var l1: u32 = 0u; - var l2: u32 = 0u; - for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - let coord = rem / out_meta.strides[d >> 2u][d & 3u]; - rem = rem % out_meta.strides[d >> 2u][d & 3u]; - l1 = l1 + min(coord, in1_meta.sizes[d >> 2u][d & 3u] - 1u) * in1_meta.strides[d >> 2u][d & 3u]; - l2 = l2 + min(coord, in2_meta.sizes[d >> 2u][d & 3u] - 1u) * in2_meta.strides[d >> 2u][d & 3u]; - } - output[idx] = floor(input1[l1] / input2[l2]); -} diff --git a/backends/webgpu/runtime/ops/linear/Linear.cpp b/backends/webgpu/runtime/ops/linear/Linear.cpp index 45cc27a515d..ba82fcc36a4 100644 --- a/backends/webgpu/runtime/ops/linear/Linear.cpp +++ b/backends/webgpu/runtime/ops/linear/Linear.cpp @@ -33,6 +33,11 @@ static_assert(sizeof(LinearParams) == 16, "LinearParams must be 16 bytes"); constexpr uint32_t kTile = 32u; +LinearParams +make_linear_params(uint32_t M, uint32_t N, uint32_t K, bool has_bias) { + return {M, N, K, has_bias ? 1u : 0u}; +} + // aten.linear (+ optional bias); shared-memory tiled GEMM. void linear_impl(WebGPUGraph& graph, const std::vector& args) { // args: [input, weight, bias?, out]; out is last. bias (arg 2) is a tensor @@ -82,11 +87,7 @@ void linear_impl(WebGPUGraph& graph, const std::vector& args) { } } - LinearParams params = {}; - params.M = M; - params.N = N; - params.K = K; - params.has_bias = has_bias ? 1u : 0u; + LinearParams params = make_linear_params(M, N, K, has_bias); // Bias binding (binding 4); a 4-byte dummy satisfies it when None // (WGSL-gated). @@ -138,7 +139,7 @@ void linear_impl(WebGPUGraph& graph, const std::vector& args) { WGPUBuffer params_buf = uniform_buffer; graph.add_tensor_resize_hook( in_id, - [in_id, out_id, M, N, K, dispatch_x, dispatch_idx, params_buf]( + [in_id, out_id, M, N, K, has_bias, dispatch_x, dispatch_idx, params_buf]( WebGPUGraph& g) { const auto& d = g.cur_dims(in_id); const uint64_t numel = utils::numel_of(d); @@ -152,10 +153,7 @@ void linear_impl(WebGPUGraph& graph, const std::vector& args) { throw std::runtime_error( "WebGPU linear: live M is 0 or exceeds the build-time max"); } - LinearParams p = {}; - p.M = m; - p.N = N; - p.K = K; + LinearParams p = make_linear_params(m, N, K, has_bias); wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); g.dispatch_at(dispatch_idx).workgroup_count_x = dispatch_x; g.dispatch_at(dispatch_idx).workgroup_count_y = diff --git a/backends/webgpu/runtime/ops/logical_and/LogicalAnd.cpp b/backends/webgpu/runtime/ops/logical_and/LogicalAnd.cpp index b6a7e12830f..358cb29fff0 100644 --- a/backends/webgpu/runtime/ops/logical_and/LogicalAnd.cpp +++ b/backends/webgpu/runtime/ops/logical_and/LogicalAnd.cpp @@ -9,7 +9,7 @@ #include #include #include -#include +#include #include diff --git a/backends/webgpu/runtime/ops/logical_and/logical_and.wgsl b/backends/webgpu/runtime/ops/logical_and/logical_and.wgsl deleted file mode 100644 index 9acb583f51c..00000000000 --- a/backends/webgpu/runtime/ops/logical_and/logical_and.wgsl +++ /dev/null @@ -1,25 +0,0 @@ -@group(0) @binding(0) var t_out: array; -@group(0) @binding(1) var t_a: array; -@group(0) @binding(2) var t_b: array; - -struct Params { - num_words: u32, - pad0: u32, - pad1: u32, - pad2: u32, -} -@group(0) @binding(3) var params: Params; - -override wg_size: u32 = 64u; - -@compute @workgroup_size(wg_size, 1, 1) -fn main( - @builtin(global_invocation_id) gid: vec3, - @builtin(num_workgroups) num_workgroups: vec3) { - // bool packed 4/word; canonical 0/1 bytes -> word-wise AND == per-byte AND. - let w = gid.x + gid.y * (num_workgroups.x * wg_size); - if (w >= params.num_words) { - return; - } - t_out[w] = t_a[w] & t_b[w]; -} diff --git a/backends/webgpu/runtime/ops/logical_and/logical_and_wgsl.h b/backends/webgpu/runtime/ops/logical_binary/logical_and_wgsl.h similarity index 96% rename from backends/webgpu/runtime/ops/logical_and/logical_and_wgsl.h rename to backends/webgpu/runtime/ops/logical_binary/logical_and_wgsl.h index 6a21a77a687..3c1f861f119 100644 --- a/backends/webgpu/runtime/ops/logical_and/logical_and_wgsl.h +++ b/backends/webgpu/runtime/ops/logical_binary/logical_and_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from logical_and.wgsl - DO NOT EDIT. +// @generated from logical_binary.wgsl - DO NOT EDIT. // wgsl-sha256: cf7c1d1dbba94e429120796c9c25a6717786cca03c08f3bd1e291d5627089c20 inline constexpr const char* kLogicalAndWGSL = R"( @group(0) @binding(0) var t_out: array; diff --git a/backends/webgpu/runtime/ops/logical_or/logical_or.wgsl b/backends/webgpu/runtime/ops/logical_binary/logical_binary.wgsl similarity index 73% rename from backends/webgpu/runtime/ops/logical_or/logical_or.wgsl rename to backends/webgpu/runtime/ops/logical_binary/logical_binary.wgsl index d7e6176ba32..fb5310cfde9 100644 --- a/backends/webgpu/runtime/ops/logical_or/logical_or.wgsl +++ b/backends/webgpu/runtime/ops/logical_binary/logical_binary.wgsl @@ -16,10 +16,13 @@ override wg_size: u32 = 64u; fn main( @builtin(global_invocation_id) gid: vec3, @builtin(num_workgroups) num_workgroups: vec3) { - // bool packed 4/word; canonical 0/1 bytes -> word-wise OR == per-byte OR. + $if OP == "&": + // bool packed 4/word; canonical 0/1 bytes -> word-wise AND == per-byte AND. + $else: + // bool packed 4/word; canonical 0/1 bytes -> word-wise OR == per-byte OR. let w = gid.x + gid.y * (num_workgroups.x * wg_size); if (w >= params.num_words) { return; } - t_out[w] = t_a[w] | t_b[w]; + t_out[w] = t_a[w] ${OP} t_b[w]; } diff --git a/backends/webgpu/runtime/ops/logical_binary/logical_binary.yaml b/backends/webgpu/runtime/ops/logical_binary/logical_binary.yaml new file mode 100644 index 00000000000..65df4465b04 --- /dev/null +++ b/backends/webgpu/runtime/ops/logical_binary/logical_binary.yaml @@ -0,0 +1,13 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +logical_binary: + parameter_names_with_default_values: + OP: "&" + shader_variants: + - NAME: logical_and + - NAME: logical_or + OP: "|" diff --git a/backends/webgpu/runtime/ops/logical_or/logical_or_wgsl.h b/backends/webgpu/runtime/ops/logical_binary/logical_or_wgsl.h similarity index 96% rename from backends/webgpu/runtime/ops/logical_or/logical_or_wgsl.h rename to backends/webgpu/runtime/ops/logical_binary/logical_or_wgsl.h index d64898cb523..e61317d2ff1 100644 --- a/backends/webgpu/runtime/ops/logical_or/logical_or_wgsl.h +++ b/backends/webgpu/runtime/ops/logical_binary/logical_or_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from logical_or.wgsl - DO NOT EDIT. +// @generated from logical_binary.wgsl - DO NOT EDIT. // wgsl-sha256: 4ad19ee04e2c7b396b4669cf44f95133d658c3ec2e6f37d7b271bedc0e582ecf inline constexpr const char* kLogicalOrWGSL = R"( @group(0) @binding(0) var t_out: array; diff --git a/backends/webgpu/runtime/ops/logical_or/LogicalOr.cpp b/backends/webgpu/runtime/ops/logical_or/LogicalOr.cpp index d26f6b486d2..d750af11227 100644 --- a/backends/webgpu/runtime/ops/logical_or/LogicalOr.cpp +++ b/backends/webgpu/runtime/ops/logical_or/LogicalOr.cpp @@ -9,7 +9,7 @@ #include #include #include -#include +#include #include diff --git a/backends/webgpu/runtime/ops/minimum/BinaryOp.cpp b/backends/webgpu/runtime/ops/minimum/BinaryOp.cpp index 3c150fc2580..f457bc9dc99 100644 --- a/backends/webgpu/runtime/ops/minimum/BinaryOp.cpp +++ b/backends/webgpu/runtime/ops/minimum/BinaryOp.cpp @@ -8,7 +8,7 @@ #include #include -#include +#include #include diff --git a/backends/webgpu/runtime/ops/minimum/binary_minimum.wgsl b/backends/webgpu/runtime/ops/minimum/binary_minimum.wgsl deleted file mode 100644 index e79cb2d2bcc..00000000000 --- a/backends/webgpu/runtime/ops/minimum/binary_minimum.wgsl +++ /dev/null @@ -1,51 +0,0 @@ -@group(0) @binding(0) var input1: array; -@group(0) @binding(1) var input2: array; -@group(0) @binding(2) var output: array; - -struct TensorMeta { - ndim: u32, - numel: u32, - sizes: array, 2>, - strides: array, 2>, -} -@group(0) @binding(3) var out_meta: TensorMeta; -@group(0) @binding(4) var in1_meta: TensorMeta; -@group(0) @binding(5) var in2_meta: TensorMeta; - -override wg_size: u32 = 64u; - -@compute @workgroup_size(wg_size, 1, 1) -fn main( - @builtin(global_invocation_id) gid: vec3, - @builtin(num_workgroups) num_workgroups: vec3) { - // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel). - let idx = gid.x + gid.y * (num_workgroups.x * wg_size); - if (idx >= out_meta.numel) { - return; - } - - // Fast path: every input dim matches the output dim -> elementwise. - var same = true; - for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - if (in1_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u] || - in2_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u]) { - same = false; - } - } - if (same) { - output[idx] = min(input1[idx], input2[idx]); - return; - } - - // Broadcast: out idx -> per-input coord (clamp size-1 dims), relinearize. - var rem = idx; - var l1: u32 = 0u; - var l2: u32 = 0u; - for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - let coord = rem / out_meta.strides[d >> 2u][d & 3u]; - rem = rem % out_meta.strides[d >> 2u][d & 3u]; - l1 = l1 + min(coord, in1_meta.sizes[d >> 2u][d & 3u] - 1u) * in1_meta.strides[d >> 2u][d & 3u]; - l2 = l2 + min(coord, in2_meta.sizes[d >> 2u][d & 3u] - 1u) * in2_meta.strides[d >> 2u][d & 3u]; - } - output[idx] = min(input1[l1], input2[l2]); -} diff --git a/backends/webgpu/runtime/ops/mul/BinaryOp.cpp b/backends/webgpu/runtime/ops/mul/BinaryOp.cpp index fdb7984b9e1..27bff334dd3 100644 --- a/backends/webgpu/runtime/ops/mul/BinaryOp.cpp +++ b/backends/webgpu/runtime/ops/mul/BinaryOp.cpp @@ -10,7 +10,7 @@ #include #include #include -#include +#include #include diff --git a/backends/webgpu/runtime/ops/mul/binary_mul.wgsl b/backends/webgpu/runtime/ops/mul/binary_mul.wgsl deleted file mode 100644 index f82a16e4b21..00000000000 --- a/backends/webgpu/runtime/ops/mul/binary_mul.wgsl +++ /dev/null @@ -1,51 +0,0 @@ -@group(0) @binding(0) var input1: array; -@group(0) @binding(1) var input2: array; -@group(0) @binding(2) var output: array; - -struct TensorMeta { - ndim: u32, - numel: u32, - sizes: array, 2>, - strides: array, 2>, -} -@group(0) @binding(3) var out_meta: TensorMeta; -@group(0) @binding(4) var in1_meta: TensorMeta; -@group(0) @binding(5) var in2_meta: TensorMeta; - -override wg_size: u32 = 64u; - -@compute @workgroup_size(wg_size, 1, 1) -fn main( - @builtin(global_invocation_id) gid: vec3, - @builtin(num_workgroups) num_workgroups: vec3) { - // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel). - let idx = gid.x + gid.y * (num_workgroups.x * wg_size); - if (idx >= out_meta.numel) { - return; - } - - // Fast path: every input dim matches the output dim -> elementwise. - var same = true; - for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - if (in1_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u] || - in2_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u]) { - same = false; - } - } - if (same) { - output[idx] = input1[idx] * input2[idx]; - return; - } - - // Broadcast: out idx -> per-input coord (clamp size-1 dims), relinearize. - var rem = idx; - var l1: u32 = 0u; - var l2: u32 = 0u; - for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - let coord = rem / out_meta.strides[d >> 2u][d & 3u]; - rem = rem % out_meta.strides[d >> 2u][d & 3u]; - l1 = l1 + min(coord, in1_meta.sizes[d >> 2u][d & 3u] - 1u) * in1_meta.strides[d >> 2u][d & 3u]; - l2 = l2 + min(coord, in2_meta.sizes[d >> 2u][d & 3u] - 1u) * in2_meta.strides[d >> 2u][d & 3u]; - } - output[idx] = input1[l1] * input2[l2]; -} diff --git a/backends/webgpu/runtime/ops/mul/silu_mul_fused.wgsl b/backends/webgpu/runtime/ops/mul/silu_mul_fused.wgsl index f8ae8136222..8281791cda6 100644 --- a/backends/webgpu/runtime/ops/mul/silu_mul_fused.wgsl +++ b/backends/webgpu/runtime/ops/mul/silu_mul_fused.wgsl @@ -7,14 +7,13 @@ struct Params { } @group(0) @binding(3) var params: Params; -// Fused SwiGLU activation: output = (g * sigmoid(g)) * up, folding the separate -// sigmoid(gate) -> mul(gate,sig)=silu -> mul(silu,up) triple into one dispatch. -// sigmoid + silu are computed in registers (never written to memory), so gate + up -// are read once and one output is written. The sigmoid form (1/(1+exp(-x))) and the -// multiply order match the original ops -> bit-exact. -@compute @workgroup_size(64) -fn main(@builtin(global_invocation_id) gid: vec3) { - let idx = gid.x; +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); if (idx >= params.num_elements) { return; } diff --git a/backends/webgpu/runtime/ops/mul/silu_mul_fused_wgsl.h b/backends/webgpu/runtime/ops/mul/silu_mul_fused_wgsl.h index e847eee4429..995ce8fe39c 100644 --- a/backends/webgpu/runtime/ops/mul/silu_mul_fused_wgsl.h +++ b/backends/webgpu/runtime/ops/mul/silu_mul_fused_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from silu_mul_fused.wgsl - DO NOT EDIT. -// wgsl-sha256: 4b8ede66c5dbc9829ff48f745eb9ad48fa5a5200058baa532fbf34f78ec2f560 +// wgsl-sha256: 7ba46c3ec15bfe4ab77a6a3e8e9f81dcb53a82328da953a3a9252fc7d470f461 inline constexpr const char* kSiluMulFusedWGSL = R"( @group(0) @binding(0) var gate: array; @group(0) @binding(1) var up: array; @@ -24,14 +24,13 @@ struct Params { } @group(0) @binding(3) var params: Params; -// Fused SwiGLU activation: output = (g * sigmoid(g)) * up, folding the separate -// sigmoid(gate) -> mul(gate,sig)=silu -> mul(silu,up) triple into one dispatch. -// sigmoid + silu are computed in registers (never written to memory), so gate + up -// are read once and one output is written. The sigmoid form (1/(1+exp(-x))) and the -// multiply order match the original ops -> bit-exact. -@compute @workgroup_size(64) -fn main(@builtin(global_invocation_id) gid: vec3) { - let idx = gid.x; +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); if (idx >= params.num_elements) { return; } diff --git a/backends/webgpu/runtime/ops/pow/BinaryOp.cpp b/backends/webgpu/runtime/ops/pow/BinaryOp.cpp index d2f923880bf..1ac90015ffe 100644 --- a/backends/webgpu/runtime/ops/pow/BinaryOp.cpp +++ b/backends/webgpu/runtime/ops/pow/BinaryOp.cpp @@ -8,7 +8,7 @@ #include #include -#include +#include #include diff --git a/backends/webgpu/runtime/ops/pow/binary_pow.wgsl b/backends/webgpu/runtime/ops/pow/binary_pow.wgsl deleted file mode 100644 index 2114ef87fee..00000000000 --- a/backends/webgpu/runtime/ops/pow/binary_pow.wgsl +++ /dev/null @@ -1,51 +0,0 @@ -@group(0) @binding(0) var input1: array; -@group(0) @binding(1) var input2: array; -@group(0) @binding(2) var output: array; - -struct TensorMeta { - ndim: u32, - numel: u32, - sizes: array, 2>, - strides: array, 2>, -} -@group(0) @binding(3) var out_meta: TensorMeta; -@group(0) @binding(4) var in1_meta: TensorMeta; -@group(0) @binding(5) var in2_meta: TensorMeta; - -override wg_size: u32 = 64u; - -@compute @workgroup_size(wg_size, 1, 1) -fn main( - @builtin(global_invocation_id) gid: vec3, - @builtin(num_workgroups) num_workgroups: vec3) { - // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel). - let idx = gid.x + gid.y * (num_workgroups.x * wg_size); - if (idx >= out_meta.numel) { - return; - } - - // Fast path: every input dim matches the output dim -> elementwise. - var same = true; - for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - if (in1_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u] || - in2_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u]) { - same = false; - } - } - if (same) { - output[idx] = pow(input1[idx], input2[idx]); - return; - } - - // Broadcast: out idx -> per-input coord (clamp size-1 dims), relinearize. - var rem = idx; - var l1: u32 = 0u; - var l2: u32 = 0u; - for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - let coord = rem / out_meta.strides[d >> 2u][d & 3u]; - rem = rem % out_meta.strides[d >> 2u][d & 3u]; - l1 = l1 + min(coord, in1_meta.sizes[d >> 2u][d & 3u] - 1u) * in1_meta.strides[d >> 2u][d & 3u]; - l2 = l2 + min(coord, in2_meta.sizes[d >> 2u][d & 3u] - 1u) * in2_meta.strides[d >> 2u][d & 3u]; - } - output[idx] = pow(input1[l1], input2[l2]); -} diff --git a/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp index 5280430df45..fe60e8377a1 100644 --- a/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp +++ b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp @@ -8,20 +8,13 @@ #include #include +#include #include #include -#include -#include -#include -#include -#include -#include -#include #include #include -#include #include #include #include @@ -60,8 +53,20 @@ constexpr uint32_t kQ4gswShmemNMinDim = 2048u; // steel GEMM: 64x64 tile, 256 threads (16x16), fixed wg (no override). constexpr uint32_t kQ4gswSteelTile = 64u; constexpr uint32_t kQ4gswSteelBK = 16u; +constexpr uint32_t kQ4gswSteelBK64 = 64u; constexpr uint32_t kQ4gswSteelInvocations = 256u; +constexpr const char* kQ4gswLinearShader = "q4gsw_linear"; +constexpr const char* kQ4gswBicolShader = "q4gsw_linear_coop4_bicol"; +constexpr const char* kQ4gswShmemShader = "q4gsw_linear_gemm_shmem"; +constexpr const char* kQ4gswSteelShader = "q4gsw_linear_gemm_steel"; +constexpr const char* kQ4gswSteelHalfShader = "q4gsw_linear_gemm_steel_half"; +constexpr const char* kQ4gswSteelHalfPwdqShader = + "q4gsw_linear_gemm_steel_half_pwdq"; +constexpr const char* kQ4gswSteelHalfPwdqF16accShader = + "q4gsw_linear_gemm_steel_half_pwdq_f16acc"; +constexpr const char* kQ4gswSteelBk64Shader = "q4gsw_steel_bk64"; + // One workgroup per (tile_m x tile_n) tile, no grid-stride: throw when the tile // count would exceed the 1D dispatch limit. Shared by the steel + shmem GEMM // routes; `kind` names the route in the error message. @@ -90,6 +95,27 @@ bool steel_supported(WGPUDevice device) { limits.maxComputeInvocationsPerWorkgroup >= kQ4gswSteelInvocations; } +bool steel_bk64_eligible( + WGPUDevice device, + uint32_t K, + uint32_t N, + uint32_t group_size, + bool has_bias) { + WGPULimits limits = {}; + if (wgpuDeviceGetLimits(device, &limits) != WGPUStatus_Success) { + return false; + } + const WebGPUContext* context = get_default_webgpu_context(); + return utils::is_q4gsw_bk64_eligible( + K, + N, + group_size, + has_bias, + context != nullptr && context->shader_f16_supported, + limits.maxComputeInvocationsPerWorkgroup, + limits.maxComputeWorkgroupStorageSize); +} + // Not grid-strided: 0 (fall back) when K%BK != 0 or over the 1D dispatch limit. uint32_t steel_workgroup_count(WGPUDevice device, uint32_t m, uint32_t n, uint32_t K) { @@ -103,6 +129,17 @@ steel_workgroup_count(WGPUDevice device, uint32_t m, uint32_t n, uint32_t K) { return (total == 0u || total > max_count) ? 0u : static_cast(total); } +uint32_t steel_bk64_workgroup_count( + WGPUDevice device, + uint32_t m, + uint32_t n, + uint32_t K) { + if (K % kQ4gswSteelBK64 != 0u) { + return 0u; + } + return steel_workgroup_count(device, m, n, K); +} + // Workgroup count for a linear_q4gsw dispatch (bicol GEMV / shmem GEMM / tiled // GEMM), with the range/limit guards shared by the build-time path and the // resize hook. use_gemv/use_shmem_gemm are the build-time routing decision (the @@ -110,10 +147,12 @@ steel_workgroup_count(WGPUDevice device, uint32_t m, uint32_t n, uint32_t K) { uint32_t compute_q4gsw_workgroup_count( WGPUDevice device, bool use_gemv, + bool use_bk64, bool use_steel, bool use_shmem_gemm, uint32_t m, uint32_t n, + uint32_t K, uint32_t wg_size, const char* op_name) { if (use_gemv) { @@ -132,6 +171,14 @@ uint32_t compute_q4gsw_workgroup_count( } return wgc; } + if (use_bk64) { + const uint32_t count = steel_bk64_workgroup_count(device, m, n, K); + if (count == 0u) { + throw std::runtime_error( + std::string("WebGPU ") + op_name + ": invalid BK64 dispatch"); + } + return count; + } if (use_steel) { // steel: one workgroup per 64x64 tile. Over-limit THROWS here -- unlike the // build-time steel_workgroup_count, which returns 0 so the caller falls @@ -162,6 +209,147 @@ uint32_t compute_q4gsw_workgroup_count( device, static_cast(total_tiles), wg_size, op_name); } +struct Q4gswExecutionState { + Q4gswParams params; + std::vector output_dims; + size_t active_route; + utils::WgCount active_grid; +}; + +constexpr size_t kQ4gswBicolRoute = 0; +constexpr size_t kQ4gswBk64Route = 1; +constexpr size_t kQ4gswPrefillRoute = 2; +// 2-route (bicol + prefill) layout, used when the BK64 route is not recorded: +// the prefill dispatch sits at index 1, not kQ4gswPrefillRoute (the 3-route 2). +constexpr size_t kQ4gswPrefillRoute2Way = 1; + +Q4gswExecutionState make_q4gsw_execution_state( + WGPUDevice device, + const std::vector& input_dims, + uint32_t max_m, + uint32_t K, + uint32_t N, + uint32_t K_packed, + uint32_t gs, + uint32_t padded_N, + uint32_t has_bias, + uint32_t wg_size, + bool use_single_gemv, + bool use_dual_route, + bool record_bk64_route, + bool bk64_eligible, + bool prefill_use_steel, + bool prefill_use_shmem_gemm) { + if (input_dims.empty()) { + throw std::runtime_error("WebGPU linear_q4gsw(resize): empty input dims"); + } + const uint64_t numel = utils::numel_of(input_dims); + if (numel % static_cast(K) != 0u) { + throw std::runtime_error( + "WebGPU linear_q4gsw(resize): live input numel not a multiple of K"); + } + const uint64_t live_m = numel / static_cast(K); + if (live_m == 0u) { + throw std::runtime_error("WebGPU linear_q4gsw(resize): live M == 0"); + } + if (live_m > max_m) { + throw std::runtime_error( + "WebGPU linear_q4gsw(resize): live M exceeds the build-time max"); + } + const uint32_t m = static_cast(live_m); + const bool use_gemv = use_single_gemv || (use_dual_route && m == 1u); + const bool use_bk64 = !use_gemv && bk64_eligible && + utils::is_q4gsw_bk64_live_m(m) && + steel_bk64_workgroup_count(device, m, N, K) > 0u; + const uint32_t workgroup_count = compute_q4gsw_workgroup_count( + device, + use_gemv, + use_bk64, + !use_gemv && !use_bk64 && prefill_use_steel, + !use_gemv && !use_bk64 && prefill_use_shmem_gemm, + m, + N, + K, + wg_size, + "linear_q4gsw(resize)"); + + Q4gswExecutionState state = {}; + state.params.M = m; + state.params.N = N; + state.params.K = K; + state.params.K_packed = K_packed; + state.params.group_size = gs; + state.params.padded_N = padded_N; + state.params.has_bias = has_bias; + state.output_dims = input_dims; + state.output_dims.back() = static_cast(N); + state.active_route = use_dual_route + ? (use_gemv ? kQ4gswBicolRoute + : (record_bk64_route + ? (use_bk64 ? kQ4gswBk64Route : kQ4gswPrefillRoute) + : kQ4gswPrefillRoute2Way)) + : 0u; + state.active_grid = {workgroup_count, 1u}; + return state; +} + +struct Q4gswResizeContext { + int in_id; + int out_id; + uint32_t max_m; + uint32_t K; + uint32_t N; + uint32_t K_packed; + uint32_t group_size; + uint32_t padded_N; + uint32_t has_bias; + uint32_t wg_size; + bool use_single_gemv; + bool use_dual_route; + bool record_bk64_route; + bool bk64_eligible; + bool prefill_use_steel; + bool prefill_use_shmem_gemm; + size_t dispatch_idx; + size_t route_group; + WGPUBuffer params_buffer; +}; + +void resize_q4gsw(WebGPUGraph& graph, const Q4gswResizeContext& context) { + const Q4gswExecutionState state = make_q4gsw_execution_state( + graph.device(), + graph.cur_dims(context.in_id), + context.max_m, + context.K, + context.N, + context.K_packed, + context.group_size, + context.padded_N, + context.has_bias, + context.wg_size, + context.use_single_gemv, + context.use_dual_route, + context.record_bk64_route, + context.bk64_eligible, + context.prefill_use_steel, + context.prefill_use_shmem_gemm); + wgpuQueueWriteBuffer( + graph.queue(), + context.params_buffer, + 0, + &state.params, + sizeof(state.params)); + if (context.use_dual_route) { + graph.select_dispatch_route( + context.route_group, state.active_route, {state.active_grid}); + } else { + auto& dispatch = graph.dispatch_at(context.dispatch_idx); + dispatch.workgroup_count_x = state.active_grid.x; + dispatch.workgroup_count_y = state.active_grid.y; + } + graph.set_cur_dims(context.out_id, state.output_dims); +} + // et_vk.linear_q4gsw args: [in, weight, scales, group_size, bias, out]. void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { const int in_id = args.at(0); @@ -240,14 +428,45 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { "WebGPU linear_q4gsw: scales dims too small for K/N"); } - // M==1 -> bicol GEMV; M>1 -> steel GEMM (preferred) else shmem else tiled. - const uint32_t wg_size = - utils::clamp_workgroup_size(device, kQ4gswLinearWorkgroupSizeX); - const bool use_gemv = (M == 1u && K % 8u == 0u && gs % 8u == 0u); + // Optional bias: real buffer if present, else a dummy for the fixed layout. + uint32_t has_bias = 0; + WGPUBuffer bias_buffer = nullptr; + uint64_t bias_size = 4; + if (graph.get_value_type(bias_id) == WebGPUGraph::ValueType::Tensor) { + const auto& bias = graph.get_tensor(bias_id); + if (bias.buffer == nullptr || bias.nbytes < N * sizeof(float)) { + throw std::runtime_error( + "WebGPU linear_q4gsw: bias present but null/undersized"); + } + has_bias = 1; + bias_buffer = bias.buffer; + bias_size = bias.nbytes; + } + if (bias_buffer == nullptr) { + bias_buffer = graph.create_scratch_buffer(4); + } + + // M==1 -> bicol GEMV; M>1 -> BK64 for exact Llama projections/M values, + // otherwise steel GEMM (preferred), shmem, or tiled. + const uint32_t wg_size = utils::clamp_workgroup_size( + device, get_webgpu_shader_info(kQ4gswLinearShader).workgroup_size_x); + const bool bicol_eligible = K % 8u == 0u && gs % 8u == 0u; + const bool use_gemv = M == 1u && bicol_eligible; + const bool use_dual_route = utils::should_record_q4gsw_dual_route( + M, + bicol_eligible, + graph.has_dynamic_shapes(), + graph.config().record_q4gsw_decode_route); + const bool bk64_eligible = + steel_bk64_eligible(device, K, N, gs, has_bias != 0u); + const bool record_bk64_route = use_dual_route && bk64_eligible && M >= 128u; + const bool use_bk64 = !use_gemv && bk64_eligible && + utils::is_q4gsw_bk64_live_m(M) && + steel_bk64_workgroup_count(device, M, N, K) > 0u; // GEMV (bicol) is a pow2 tree reduction; compute its size only when used. - const uint32_t gemv_wg_size = use_gemv + const uint32_t gemv_wg_size = (use_gemv || use_dual_route) ? utils::clamp_workgroup_size_pow2( - device, kQ4gswLinearCoop4BicolWorkgroupSizeX) + device, get_webgpu_shader_info(kQ4gswBicolShader).workgroup_size_x) : 0u; // steel (256-thread) is the preferred M>1 prefill GEMM; 0 count = ineligible. const bool use_steel = !use_gemv && steel_supported(device) && @@ -258,10 +477,14 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { // large K/N thresholds; otherwise the register-tiled path handles it. const bool use_shmem_gemm = !use_gemv && !use_steel && (K >= kQ4gswShmemMinDim || N >= kQ4gswShmemNMinDim); - const char* shader_src = use_gemv ? kQ4gswLinearCoop4BicolWGSL - : use_steel ? kQ4gswLinearGemmSteelWGSL - : use_shmem_gemm ? kQ4gswLinearGemmShmemWGSL - : kQ4gswLinearWGSL; + const char* prefill_shader_name = use_steel ? kQ4gswSteelShader + : use_shmem_gemm ? kQ4gswShmemShader + : kQ4gswLinearShader; + const char* shader_name = use_gemv ? kQ4gswBicolShader + : use_bk64 ? kQ4gswSteelBk64Shader + : use_steel ? kQ4gswSteelShader + : use_shmem_gemm ? kQ4gswShmemShader + : kQ4gswLinearShader; // f16-multiply steel: only when the device negotiated shader-f16; else the // f32 steel kernel runs (fail-closed). Same bindings and tile. if (use_steel) { @@ -271,9 +494,12 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { // each u32 weight word once + hoists the per-column scale (half re-reads // them ~8x/~16x). Needs group_size % BK == 0 so the hoisted scale is // constant across the BK tile; else the per-nibble `half` kernel. - shader_src = (gs % kQ4gswSteelBK == 0u) - ? kQ4gswLinearGemmSteelHalfPwdqWGSL - : kQ4gswLinearGemmSteelHalfWGSL; + prefill_shader_name = (gs % kQ4gswSteelBK == 0u) + ? kQ4gswSteelHalfPwdqShader + : kQ4gswSteelHalfShader; + if (!use_bk64) { + shader_name = prefill_shader_name; + } } } // f16-accumulate: pwdq staging with an f16 register accumulator. @@ -284,158 +510,157 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { if (use_steel && graph.f16_accumulate_gemm() && (gs % kQ4gswSteelBK == 0u)) { const WebGPUContext* ctx = get_default_webgpu_context(); if (ctx != nullptr && ctx->shader_f16_supported) { - shader_src = kQ4gswLinearGemmSteelHalfPwdqF16accWGSL; + prefill_shader_name = kQ4gswSteelHalfPwdqF16accShader; + if (!use_bk64) { + shader_name = prefill_shader_name; + } } } - const uint32_t workgroup_count = compute_q4gsw_workgroup_count( + + const Q4gswExecutionState initial_state = make_q4gsw_execution_state( device, - use_gemv, - use_steel, - use_shmem_gemm, + in.dims, M, + K, N, + K_packed, + gs, + padded_N, + has_bias, wg_size, - "linear_q4gsw"); - - // Optional bias: real buffer if present, else a dummy for the fixed layout. - uint32_t has_bias = 0; - WGPUBuffer bias_buffer = nullptr; - uint64_t bias_size = 4; - if (graph.get_value_type(bias_id) == WebGPUGraph::ValueType::Tensor) { - const auto& bias = graph.get_tensor(bias_id); - if (bias.buffer == nullptr || bias.nbytes < N * sizeof(float)) { - throw std::runtime_error( - "WebGPU linear_q4gsw: bias present but null/undersized"); + use_gemv, + use_dual_route, + record_bk64_route, + bk64_eligible, + use_steel, + use_shmem_gemm); + + WGPUBuffer params_buffer = graph.create_params_buffer(initial_state.params); + const std::vector bindings = { + {0, WGPUBufferBindingType_Storage, out.buffer, out.nbytes}, + {1, WGPUBufferBindingType_ReadOnlyStorage, in.buffer, in.nbytes}, + {2, WGPUBufferBindingType_ReadOnlyStorage, weight.buffer, weight.nbytes}, + {3, WGPUBufferBindingType_ReadOnlyStorage, scales.buffer, scales.nbytes}, + {4, WGPUBufferBindingType_ReadOnlyStorage, bias_buffer, bias_size}, + {5, WGPUBufferBindingType_Uniform, params_buffer, sizeof(Q4gswParams)}, + }; + auto make_bundle = + [&](const char* source, bool fixed_wg, uint32_t override_wg_size) { + const WGPUConstantEntry wg_size_constant = + utils::make_wg_size_constant(override_wg_size); + return utils::make_compute_pipeline( + device, + source, + bindings, + fixed_wg ? nullptr : &wg_size_constant, + fixed_wg ? 0u : 1u); + }; + auto make_shared_bundle = + [&](const char* source, + const utils::ComputePipelineBundle& shared_resources, + bool fixed_wg, + uint32_t override_wg_size) { + const WGPUConstantEntry wg_size_constant = + utils::make_wg_size_constant(override_wg_size); + return utils::make_compute_pipeline( + device, + source, + shared_resources, + fixed_wg ? nullptr : &wg_size_constant, + fixed_wg ? 0u : 1u); + }; + + const bool fixed_prefill_wg = use_steel || use_shmem_gemm; + const char* prefill_label = use_steel ? "linear_q4gsw_steel" + : use_shmem_gemm ? "linear_q4gsw_shmem" + : "linear_q4gsw_tiled"; + size_t dispatch_idx = 0; + size_t route_group = 0; + if (use_dual_route) { + utils::ComputePipelineBundle bicol_bundle = make_bundle( + get_webgpu_shader_info(kQ4gswBicolShader).source, false, gemv_wg_size); + const size_t bicol_idx = graph.add_dispatch( + {bicol_bundle.pipeline, + bicol_bundle.bind_group, + initial_state.active_grid.x, + "linear_q4gsw_coop4_bicol", + initial_state.active_grid.y}); + size_t bk64_idx = 0; + if (record_bk64_route) { + utils::ComputePipelineBundle bk64_bundle = make_shared_bundle( + get_webgpu_shader_info(kQ4gswSteelBk64Shader).source, + bicol_bundle, + true, + 0u); + bk64_idx = graph.add_dispatch( + {bk64_bundle.pipeline, + bk64_bundle.bind_group, + initial_state.active_grid.x, + "linear_q4gsw_bk64", + initial_state.active_grid.y}); } - has_bias = 1; - bias_buffer = bias.buffer; - bias_size = bias.nbytes; - } - if (bias_buffer == nullptr) { - bias_buffer = graph.create_scratch_buffer(4); + utils::ComputePipelineBundle prefill_bundle = make_shared_bundle( + get_webgpu_shader_info(prefill_shader_name).source, + bicol_bundle, + fixed_prefill_wg, + wg_size); + const size_t prefill_idx = graph.add_dispatch( + {prefill_bundle.pipeline, + prefill_bundle.bind_group, + initial_state.active_grid.x, + prefill_label, + initial_state.active_grid.y}); + if (record_bk64_route) { + route_group = graph.register_dispatch_route_group( + {{bicol_idx, bicol_idx + 1}, + {bk64_idx, bk64_idx + 1}, + {prefill_idx, prefill_idx + 1}}); + } else { + route_group = graph.register_dispatch_route_group( + {{bicol_idx, bicol_idx + 1}, {prefill_idx, prefill_idx + 1}}); + } + graph.select_dispatch_route( + route_group, initial_state.active_route, {initial_state.active_grid}); + } else { + const bool fixed_wg = use_gemv ? false : (use_bk64 || fixed_prefill_wg); + utils::ComputePipelineBundle bundle = make_bundle( + get_webgpu_shader_info(shader_name).source, + fixed_wg, + use_gemv ? gemv_wg_size : wg_size); + dispatch_idx = graph.add_dispatch( + {bundle.pipeline, + bundle.bind_group, + initial_state.active_grid.x, + use_gemv ? "linear_q4gsw_coop4_bicol" + : (use_bk64 ? "linear_q4gsw_bk64" : prefill_label), + initial_state.active_grid.y}); } - Q4gswParams params = {}; - params.M = M; - params.N = N; - params.K = K; - params.K_packed = K_packed; - params.group_size = gs; - params.padded_N = padded_N; - params.has_bias = has_bias; - - WGPUBufferDescriptor uniform_desc = {}; - uniform_desc.size = sizeof(Q4gswParams); - uniform_desc.usage = WGPUBufferUsage_Uniform | WGPUBufferUsage_CopyDst; - uniform_desc.mappedAtCreation = true; - WGPUBuffer uniform_buffer = wgpuDeviceCreateBuffer(device, &uniform_desc); - void* mapped = - wgpuBufferGetMappedRange(uniform_buffer, 0, sizeof(Q4gswParams)); - std::memcpy(mapped, ¶ms, sizeof(Q4gswParams)); - wgpuBufferUnmap(uniform_buffer); - graph.add_uniform_buffer_bytes(sizeof(Q4gswParams)); - - // GEMV/tiled wire an override wg_size; steel (256) + shmem (64) are fixed. - const bool fixed_wg = use_steel || use_shmem_gemm; - WGPUConstantEntry wg_size_constant = {}; - wg_size_constant.key = {"wg_size", WGPU_STRLEN}; - wg_size_constant.value = - static_cast(use_gemv ? gemv_wg_size : wg_size); - - utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( - device, - shader_src, - { - {0, WGPUBufferBindingType_Storage, out.buffer, out.nbytes}, - {1, WGPUBufferBindingType_ReadOnlyStorage, in.buffer, in.nbytes}, - {2, - WGPUBufferBindingType_ReadOnlyStorage, - weight.buffer, - weight.nbytes}, - {3, - WGPUBufferBindingType_ReadOnlyStorage, - scales.buffer, - scales.nbytes}, - {4, WGPUBufferBindingType_ReadOnlyStorage, bias_buffer, bias_size}, - {5, - WGPUBufferBindingType_Uniform, - uniform_buffer, - sizeof(Q4gswParams)}, - }, - fixed_wg ? nullptr : &wg_size_constant, - fixed_wg ? 0u : 1u); - - const size_t dispatch_idx = graph.add_dispatch( - {bundle.pipeline, bundle.bind_group, workgroup_count, "linear_q4gsw"}); - - // Dynamic shapes: recompute dispatch + params.M for the live M. use_gemv and - // use_shmem_gemm are captured (routing is fixed at build); the helper re-runs - // the same path's workgroup-count formula with the live m. - graph.add_tensor_resize_hook( + // Dynamic shapes: recompute one shared Params block and select exactly one + // writer. The prefill pipeline remains the route chosen from max M. + const Q4gswResizeContext resize_context = { in_id, - [in_id, - out_id, - M, - K, - N, - K_packed, - gs, - padded_N, - has_bias, - wg_size, - use_gemv, - use_steel, - use_shmem_gemm, - dispatch_idx, - uniform_buffer](WebGPUGraph& g) { - const auto& d = g.cur_dims(in_id); - if (d.empty()) { - throw std::runtime_error( - "WebGPU linear_q4gsw(resize): empty input dims"); - } - const uint64_t numel = utils::numel_of(d); - if (numel % static_cast(K) != 0u) { - throw std::runtime_error( - "WebGPU linear_q4gsw(resize): live input numel not a multiple " - "of K"); - } - const uint32_t m = - static_cast(numel / static_cast(K)); - if (m == 0u) { - throw std::runtime_error("WebGPU linear_q4gsw(resize): live M == 0"); - } - // Buffers/bind-groups were sized for the build-time max M; a larger - // live M would write out of bounds. - if (m > M) { - throw std::runtime_error( - "WebGPU linear_q4gsw(resize): live M exceeds the build-time max"); - } - const uint32_t wgc = compute_q4gsw_workgroup_count( - g.device(), - use_gemv, - use_steel, - use_shmem_gemm, - m, - N, - wg_size, - "linear_q4gsw(resize)"); - Q4gswParams p = {}; - p.M = m; - p.N = N; - p.K = K; - p.K_packed = K_packed; - p.group_size = gs; - p.padded_N = padded_N; - p.has_bias = has_bias; - wgpuQueueWriteBuffer(g.queue(), uniform_buffer, 0, &p, sizeof(p)); - g.dispatch_at(dispatch_idx).workgroup_count_x = wgc; - std::vector od(d.begin(), d.end()); - od.back() = static_cast(N); - g.set_cur_dims(out_id, od); - }); - - // Graph owns it so the resize hook can rewrite it; freed in the dtor. - graph.own_uniform_buffer(uniform_buffer); + out_id, + M, + K, + N, + K_packed, + gs, + padded_N, + has_bias, + wg_size, + use_gemv, + use_dual_route, + record_bk64_route, + bk64_eligible, + use_steel, + use_shmem_gemm, + dispatch_idx, + route_group, + params_buffer}; + graph.add_tensor_resize_hook(in_id, [resize_context](WebGPUGraph& g) { + resize_q4gsw(g, resize_context); + }); } } // namespace diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel.wgsl b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel.wgsl index e2dfc610976..19d2b9aff77 100644 --- a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel.wgsl +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel.wgsl @@ -1,6 +1,6 @@ $if DTYPE == "half": enable f16; -@group(0) @binding(0) var t_out: array; +@group(0) @binding(0) var t_out: array<${"vec4" if BK == 64 else "f32"}>; @group(0) @binding(1) var t_input: array>; @group(0) @binding(2) var t_weight: array; @group(0) @binding(3) var t_scales: array; @@ -31,9 +31,10 @@ struct Params { // ACC=half (PWDQ only) f16 accumulate with fma(), cast to f32 in the epilogue // -- LOSSY, perplexity-gated, opt-in via a runtime spec. ACC=float is f32 // accumulate -- BIT-EXACT to the per-nibble half kernel. -const BM: u32 = 64u; const BN: u32 = 64u; const BK: u32 = 16u; -var As: array<${buffer_scalar_type(DTYPE)}, 1024>; // BM*BK -var Bs: array<${buffer_scalar_type(DTYPE)}, 1024>; // BK*BN +// BK=64 (PWDQ + ACC=half only) stages a full quantization group at once. +const BM: u32 = 64u; const BN: u32 = 64u; const BK: u32 = ${BK}u; +var As: array<${buffer_scalar_type(DTYPE)}, ${64 * BK}>; // BM*BK +var Bs: array<${buffer_scalar_type(DTYPE)}, ${BK * 64}>; // BK*BN // 16x16 = 256 threads, bound to the 64x64 tile + 4x4 reg tile (not a knob). @compute @workgroup_size(16, 16) fn main(@builtin(workgroup_id) wid: vec3, @@ -64,14 +65,35 @@ fn main(@builtin(workgroup_id) wid: vec3, if (arow < params.M) { let base = arow * params.K + k0 + ac; // vec4 coalesced load; base is 4-aligned on the steel route (K%16==0, ac/k0 multiples of 4). - let av = t_input[base >> 2u]; - As[ar * BK + ac + 0u] = ${buffer_scalar_type(DTYPE)}(av.x); - As[ar * BK + ac + 1u] = ${buffer_scalar_type(DTYPE)}(av.y); - As[ar * BK + ac + 2u] = ${buffer_scalar_type(DTYPE)}(av.z); - As[ar * BK + ac + 3u] = ${buffer_scalar_type(DTYPE)}(av.w); + $if BK == 64: + let av0 = t_input[(base + 0u) >> 2u]; + let av1 = t_input[(base + 16u) >> 2u]; + let av2 = t_input[(base + 32u) >> 2u]; + let av3 = t_input[(base + 48u) >> 2u]; + As[ar * BK + ac + 0u] = f16(av0.x); As[ar * BK + ac + 1u] = f16(av0.y); + As[ar * BK + ac + 2u] = f16(av0.z); As[ar * BK + ac + 3u] = f16(av0.w); + As[ar * BK + ac + 16u] = f16(av1.x); As[ar * BK + ac + 17u] = f16(av1.y); + As[ar * BK + ac + 18u] = f16(av1.z); As[ar * BK + ac + 19u] = f16(av1.w); + As[ar * BK + ac + 32u] = f16(av2.x); As[ar * BK + ac + 33u] = f16(av2.y); + As[ar * BK + ac + 34u] = f16(av2.z); As[ar * BK + ac + 35u] = f16(av2.w); + As[ar * BK + ac + 48u] = f16(av3.x); As[ar * BK + ac + 49u] = f16(av3.y); + As[ar * BK + ac + 50u] = f16(av3.z); As[ar * BK + ac + 51u] = f16(av3.w); + $else: + let av = t_input[base >> 2u]; + As[ar * BK + ac + 0u] = ${buffer_scalar_type(DTYPE)}(av.x); + As[ar * BK + ac + 1u] = ${buffer_scalar_type(DTYPE)}(av.y); + As[ar * BK + ac + 2u] = ${buffer_scalar_type(DTYPE)}(av.z); + As[ar * BK + ac + 3u] = ${buffer_scalar_type(DTYPE)}(av.w); } else { - As[ar * BK + ac + 0u] = ${"0.0h" if PWDQ else "0.0"}; As[ar * BK + ac + 1u] = ${"0.0h" if PWDQ else "0.0"}; - As[ar * BK + ac + 2u] = ${"0.0h" if PWDQ else "0.0"}; As[ar * BK + ac + 3u] = ${"0.0h" if PWDQ else "0.0"}; + $if BK == 64: + for (var segment: u32 = 0u; segment < 4u; segment = segment + 1u) { + for (var ai: u32 = 0u; ai < 4u; ai = ai + 1u) { + As[ar * BK + ac + segment * 16u + ai] = 0.0h; + } + } + $else: + As[ar * BK + ac + 0u] = ${"0.0h" if PWDQ else "0.0"}; As[ar * BK + ac + 1u] = ${"0.0h" if PWDQ else "0.0"}; + As[ar * BK + ac + 2u] = ${"0.0h" if PWDQ else "0.0"}; As[ar * BK + ac + 3u] = ${"0.0h" if PWDQ else "0.0"}; } $if PWDQ: // Packed-word dequant: threads [0,BN) each stage one full BK-column of Bs. @@ -83,16 +105,28 @@ fn main(@builtin(workgroup_id) wid: vec3, // group sizes; K%BK==0 on the steel route), so hoist it to one read. let scale_row = (k0 / params.group_size) * params.padded_N; let scale = f16(t_scales[scale_row + n]); - // Column n's 16-nibble K-slice for this tile = two consecutive words. + // Column n's BK-nibble K-slice starts at this packed word. // K_packed multiple of 8 => base_word stays inside column n's own region. let base_word = n * (params.K_packed >> 2u) + (k0 >> 3u); - let w0 = t_weight[base_word]; - let w1 = t_weight[base_word + 1u]; - for (var br: u32 = 0u; br < BK; br = br + 1u) { - let word = select(w1, w0, br < 8u); // word0 holds K-slice [0,8) - let nib = (word >> ((br & 7u) * 4u)) & 0x0Fu; - Bs[br * BN + c] = f16(i32(nib) - 8) * scale; - } + $if BK == 64: + let words = array( + t_weight[base_word + 0u], t_weight[base_word + 1u], + t_weight[base_word + 2u], t_weight[base_word + 3u], + t_weight[base_word + 4u], t_weight[base_word + 5u], + t_weight[base_word + 6u], t_weight[base_word + 7u]); + for (var br: u32 = 0u; br < BK; br = br + 1u) { + let word = words[br >> 3u]; + let nib = (word >> ((br & 7u) * 4u)) & 0x0Fu; + Bs[br * BN + c] = f16(i32(nib) - 8) * scale; + } + $else: + let w0 = t_weight[base_word]; + let w1 = t_weight[base_word + 1u]; + for (var br: u32 = 0u; br < BK; br = br + 1u) { + let word = select(w1, w0, br < 8u); // word0 holds K-slice [0,8) + let nib = (word >> ((br & 7u) * 4u)) & 0x0Fu; + Bs[br * BN + c] = f16(i32(nib) - 8) * scale; + } } else { for (var br: u32 = 0u; br < BK; br = br + 1u) { Bs[br * BN + c] = 0.0h; } } @@ -135,15 +169,30 @@ fn main(@builtin(workgroup_id) wid: vec3, workgroupBarrier(); k0 = k0 + BK; } - for (var m: u32 = 0u; m < 4u; m = m + 1u) { - for (var n: u32 = 0u; n < 4u; n = n + 1u) { + $if BK == 64: + for (var m: u32 = 0u; m < 4u; m = m + 1u) { let r = row0 + lid.y * 4u + m; - let c = col0 + lid.x * 4u + n; - if (r < params.M && c < params.N) { - var v = ${"f32(acc[m][n])" if ACC == "half" else "acc[m][n]"}; - if (params.has_bias != 0u) { v = v + t_bias[c]; } - t_out[r * params.N + c] = v; + let c0 = col0 + lid.x * 4u; + if (r < params.M && c0 < params.N) { + var vv = vec4( + f32(acc[m][0]), f32(acc[m][1]), f32(acc[m][2]), f32(acc[m][3])); + if (params.has_bias != 0u) { + vv = vv + vec4( + t_bias[c0], t_bias[c0 + 1u], t_bias[c0 + 2u], t_bias[c0 + 3u]); + } + t_out[(r * params.N + c0) >> 2u] = vv; + } + } + $else: + for (var m: u32 = 0u; m < 4u; m = m + 1u) { + for (var n: u32 = 0u; n < 4u; n = n + 1u) { + let r = row0 + lid.y * 4u + m; + let c = col0 + lid.x * 4u + n; + if (r < params.M && c < params.N) { + var v = ${"f32(acc[m][n])" if ACC == "half" else "acc[m][n]"}; + if (params.has_bias != 0u) { v = v + t_bias[c]; } + t_out[r * params.N + c] = v; + } } } - } } diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel.yaml b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel.yaml index 5a2cae5e499..9c16e924a0c 100644 --- a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel.yaml +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel.yaml @@ -3,6 +3,7 @@ q4gsw_linear_gemm_steel: DTYPE: float PWDQ: false ACC: float + BK: 16 shader_variants: - NAME: q4gsw_linear_gemm_steel DTYPE: float @@ -20,3 +21,8 @@ q4gsw_linear_gemm_steel: DTYPE: half PWDQ: true ACC: half + - NAME: q4gsw_steel_bk64 + DTYPE: half + PWDQ: true + ACC: half + BK: 64 diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_f16acc_wgsl.h b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_f16acc_wgsl.h index ab4d3c06915..a57292ac0f1 100644 --- a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_f16acc_wgsl.h +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_f16acc_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from q4gsw_linear_gemm_steel.wgsl - DO NOT EDIT. -// wgsl-sha256: a4346ddef028036f29aa73f0620c586467627626a745159fccef816a32f9475a +// wgsl-sha256: 99b460d282c04a4f624e03f8216d3f36ba5f5562b6f60ff752ed9a7f4becffdc inline constexpr const char* kQ4gswLinearGemmSteelHalfPwdqF16accWGSL = R"( enable f16; @group(0) @binding(0) var t_out: array; @@ -47,6 +47,7 @@ struct Params { // ACC=half (PWDQ only) f16 accumulate with fma(), cast to f32 in the epilogue // -- LOSSY, perplexity-gated, opt-in via a runtime spec. ACC=float is f32 // accumulate -- BIT-EXACT to the per-nibble half kernel. +// BK=64 (PWDQ + ACC=half only) stages a full quantization group at once. const BM: u32 = 64u; const BN: u32 = 64u; const BK: u32 = 16u; var As: array; // BM*BK var Bs: array; // BK*BN @@ -94,7 +95,7 @@ fn main(@builtin(workgroup_id) wid: vec3, // group sizes; K%BK==0 on the steel route), so hoist it to one read. let scale_row = (k0 / params.group_size) * params.padded_N; let scale = f16(t_scales[scale_row + n]); - // Column n's 16-nibble K-slice for this tile = two consecutive words. + // Column n's BK-nibble K-slice starts at this packed word. // K_packed multiple of 8 => base_word stays inside column n's own region. let base_word = n * (params.K_packed >> 2u) + (k0 >> 3u); let w0 = t_weight[base_word]; diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_wgsl.h b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_wgsl.h index 251a78b9fc8..2cb01ef7b53 100644 --- a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_wgsl.h +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from q4gsw_linear_gemm_steel.wgsl - DO NOT EDIT. -// wgsl-sha256: d772a385edf91547a51fa3f5c4bdb4101da933f4c47828578c2818ff9eadf5fc +// wgsl-sha256: ea3edb89ec54946b392d5b4f09dd027ff19b3baad8352f1da567796e4eabc047 inline constexpr const char* kQ4gswLinearGemmSteelHalfPwdqWGSL = R"( enable f16; @group(0) @binding(0) var t_out: array; @@ -47,6 +47,7 @@ struct Params { // ACC=half (PWDQ only) f16 accumulate with fma(), cast to f32 in the epilogue // -- LOSSY, perplexity-gated, opt-in via a runtime spec. ACC=float is f32 // accumulate -- BIT-EXACT to the per-nibble half kernel. +// BK=64 (PWDQ + ACC=half only) stages a full quantization group at once. const BM: u32 = 64u; const BN: u32 = 64u; const BK: u32 = 16u; var As: array; // BM*BK var Bs: array; // BK*BN @@ -94,7 +95,7 @@ fn main(@builtin(workgroup_id) wid: vec3, // group sizes; K%BK==0 on the steel route), so hoist it to one read. let scale_row = (k0 / params.group_size) * params.padded_N; let scale = f16(t_scales[scale_row + n]); - // Column n's 16-nibble K-slice for this tile = two consecutive words. + // Column n's BK-nibble K-slice starts at this packed word. // K_packed multiple of 8 => base_word stays inside column n's own region. let base_word = n * (params.K_packed >> 2u) + (k0 >> 3u); let w0 = t_weight[base_word]; diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_wgsl.h b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_wgsl.h index 4741e13b74a..37b0b8b471f 100644 --- a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_wgsl.h +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from q4gsw_linear_gemm_steel.wgsl - DO NOT EDIT. -// wgsl-sha256: 363916e1aacae9635b4573d2ce35d0ee22916ce225537a3717df98a0ed3353da +// wgsl-sha256: 95cd59521a07d686a0eb396f79fe4e260e16e81d61f3c185ed7ed9be72963d59 inline constexpr const char* kQ4gswLinearGemmSteelHalfWGSL = R"( enable f16; @group(0) @binding(0) var t_out: array; @@ -47,6 +47,7 @@ struct Params { // ACC=half (PWDQ only) f16 accumulate with fma(), cast to f32 in the epilogue // -- LOSSY, perplexity-gated, opt-in via a runtime spec. ACC=float is f32 // accumulate -- BIT-EXACT to the per-nibble half kernel. +// BK=64 (PWDQ + ACC=half only) stages a full quantization group at once. const BM: u32 = 64u; const BN: u32 = 64u; const BK: u32 = 16u; var As: array; // BM*BK var Bs: array; // BK*BN diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_wgsl.h b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_wgsl.h index 197b4f0ae92..04db8ea7a5b 100644 --- a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_wgsl.h +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from q4gsw_linear_gemm_steel.wgsl - DO NOT EDIT. -// wgsl-sha256: f8106cb3f18424a9db8e04464e52a7fd522ca78f49e19ba6692f1c8fa83474e6 +// wgsl-sha256: 11de59057b19656a2b0ab6bf8c78c37d56997f2a3664634b4a1adf07bddc45ee inline constexpr const char* kQ4gswLinearGemmSteelWGSL = R"( @group(0) @binding(0) var t_out: array; @group(0) @binding(1) var t_input: array>; @@ -46,6 +46,7 @@ struct Params { // ACC=half (PWDQ only) f16 accumulate with fma(), cast to f32 in the epilogue // -- LOSSY, perplexity-gated, opt-in via a runtime spec. ACC=float is f32 // accumulate -- BIT-EXACT to the per-nibble half kernel. +// BK=64 (PWDQ + ACC=half only) stages a full quantization group at once. const BM: u32 = 64u; const BN: u32 = 64u; const BK: u32 = 16u; var As: array; // BM*BK var Bs: array; // BK*BN diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_qkv_fused.wgsl b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_qkv_bk64.wgsl similarity index 63% rename from backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_qkv_fused.wgsl rename to backends/webgpu/runtime/ops/quantized_linear/q4gsw_qkv_bk64.wgsl index c89924a67c1..aa977b65284 100644 --- a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_qkv_fused.wgsl +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_qkv_bk64.wgsl @@ -1,12 +1,5 @@ enable f16; -// Fused QKV q4gsw GEMM (Llama attention projections): one [M, N=3072] pwdq + f16-accumulate GEMM -// (vec4 activation load) that scatter-writes each output column range to a SEPARATE buffer -- -// c<2048 -> q, [2048,2560) -> k, [2560,3072) -> v. Replaces the 3 separate q/k/v linear dispatches; -// fixes the N=512 K/V occupancy starvation (16 WGs -> 96 WGs at M~128). Boundaries are 64-tile-aligned -// so each 64-col tile maps to exactly one output (uniform branch per workgroup). Per-output ROW STRIDE: -// q=2048, k=v=512. BIT-EXACT to 3 separate pwdqf16acc linears (fusing along N does not change the -// per-column K-accumulation order). Validated on Canary M4 Pro: correct (maxRel ~1e-3), scatter overhead -// 1.02x (free), concat win 1.63x on the QKV block. Boundaries hardcoded for Llama-3.2-1B GQA (32Q/8KV). + @group(0) @binding(0) var t_out_q: array; @group(0) @binding(1) var t_out_k: array; @group(0) @binding(2) var t_out_v: array; @@ -25,10 +18,12 @@ struct Params { _pad: u32, } @group(0) @binding(7) var params: Params; -const BM: u32 = 64u; const BN: u32 = 64u; const BK: u32 = 16u; + +// BK64 QKV variant: group_size=64 keeps one scale valid for all eight packed words. +const BM: u32 = 64u; const BN: u32 = 64u; const BK: u32 = 64u; const N_Q: u32 = 2048u; const N_QK: u32 = 2560u; const N_KV: u32 = 512u; -var As: array; -var Bs: array; +var As: array; +var Bs: array; @compute @workgroup_size(16, 16) fn main(@builtin(workgroup_id) wid: vec3, @builtin(local_invocation_id) lid: vec3) { @@ -44,18 +39,31 @@ fn main(@builtin(workgroup_id) wid: vec3, } let ar = tid / 4u; let ac = (tid % 4u) * 4u; + var k0: u32 = 0u; loop { if (k0 >= params.K) { break; } let arow = row0 + ar; if (arow < params.M) { let base = arow * params.K + k0 + ac; - let av = t_input[base >> 2u]; - As[ar * BK + ac + 0u] = f16(av.x); As[ar * BK + ac + 1u] = f16(av.y); - As[ar * BK + ac + 2u] = f16(av.z); As[ar * BK + ac + 3u] = f16(av.w); + let av0 = t_input[(base + 0u) >> 2u]; + let av1 = t_input[(base + 16u) >> 2u]; + let av2 = t_input[(base + 32u) >> 2u]; + let av3 = t_input[(base + 48u) >> 2u]; + As[ar * BK + ac + 0u] = f16(av0.x); As[ar * BK + ac + 1u] = f16(av0.y); + As[ar * BK + ac + 2u] = f16(av0.z); As[ar * BK + ac + 3u] = f16(av0.w); + As[ar * BK + ac + 16u] = f16(av1.x); As[ar * BK + ac + 17u] = f16(av1.y); + As[ar * BK + ac + 18u] = f16(av1.z); As[ar * BK + ac + 19u] = f16(av1.w); + As[ar * BK + ac + 32u] = f16(av2.x); As[ar * BK + ac + 33u] = f16(av2.y); + As[ar * BK + ac + 34u] = f16(av2.z); As[ar * BK + ac + 35u] = f16(av2.w); + As[ar * BK + ac + 48u] = f16(av3.x); As[ar * BK + ac + 49u] = f16(av3.y); + As[ar * BK + ac + 50u] = f16(av3.z); As[ar * BK + ac + 51u] = f16(av3.w); } else { - As[ar * BK + ac + 0u] = 0.0h; As[ar * BK + ac + 1u] = 0.0h; - As[ar * BK + ac + 2u] = 0.0h; As[ar * BK + ac + 3u] = 0.0h; + for (var segment: u32 = 0u; segment < 4u; segment = segment + 1u) { + for (var ai: u32 = 0u; ai < 4u; ai = ai + 1u) { + As[ar * BK + ac + segment * 16u + ai] = 0.0h; + } + } } if (tid < BN) { let c = tid; @@ -64,10 +72,17 @@ fn main(@builtin(workgroup_id) wid: vec3, let scale_row = (k0 / params.group_size) * params.padded_N; let scale = f16(t_scales[scale_row + n]); let base_word = n * (params.K_packed >> 2u) + (k0 >> 3u); - let w0 = t_weight[base_word]; + let w0 = t_weight[base_word + 0u]; let w1 = t_weight[base_word + 1u]; + let w2 = t_weight[base_word + 2u]; + let w3 = t_weight[base_word + 3u]; + let w4 = t_weight[base_word + 4u]; + let w5 = t_weight[base_word + 5u]; + let w6 = t_weight[base_word + 6u]; + let w7 = t_weight[base_word + 7u]; + let words = array(w0, w1, w2, w3, w4, w5, w6, w7); for (var br: u32 = 0u; br < BK; br = br + 1u) { - let word = select(w1, w0, br < 8u); + let word = words[br >> 3u]; let nib = (word >> ((br & 7u) * 4u)) & 0x0Fu; Bs[br * BN + c] = f16(i32(nib) - 8) * scale; } @@ -91,7 +106,7 @@ fn main(@builtin(workgroup_id) wid: vec3, for (var m: u32 = 0u; m < 4u; m = m + 1u) { for (var n: u32 = 0u; n < 4u; n = n + 1u) { let r = row0 + lid.y * 4u + m; - let c = col0 + lid.x * 4u + n; // global fused column [0, 3072) + let c = col0 + lid.x * 4u + n; if (r < params.M && c < params.N) { var val = f32(acc[m][n]); if (params.has_bias != 0u) { val = val + t_bias[c]; } diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_qkv_fused_wgsl.h b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_qkv_bk64_wgsl.h similarity index 61% rename from backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_qkv_fused_wgsl.h rename to backends/webgpu/runtime/ops/quantized_linear/q4gsw_qkv_bk64_wgsl.h index 93243698a3b..371f79785ab 100644 --- a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_qkv_fused_wgsl.h +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_qkv_bk64_wgsl.h @@ -12,18 +12,11 @@ namespace executorch::backends::webgpu { -// @generated from q4gsw_linear_gemm_qkv_fused.wgsl - DO NOT EDIT. -// wgsl-sha256: 93e127e8ee4609d846015c8b75a600a29502e19a92bdf3a08e3429635f834085 -inline constexpr const char* kQ4gswLinearGemmQkvFusedWGSL = R"( +// @generated from q4gsw_qkv_bk64.wgsl - DO NOT EDIT. +// wgsl-sha256: d738762f00f79ca16cf1549d47e6d1f51155f50805eec5e7e6df3bc07ee309ee +inline constexpr const char* kQ4gswQkvBk64WGSL = R"( enable f16; -// Fused QKV q4gsw GEMM (Llama attention projections): one [M, N=3072] pwdq + f16-accumulate GEMM -// (vec4 activation load) that scatter-writes each output column range to a SEPARATE buffer -- -// c<2048 -> q, [2048,2560) -> k, [2560,3072) -> v. Replaces the 3 separate q/k/v linear dispatches; -// fixes the N=512 K/V occupancy starvation (16 WGs -> 96 WGs at M~128). Boundaries are 64-tile-aligned -// so each 64-col tile maps to exactly one output (uniform branch per workgroup). Per-output ROW STRIDE: -// q=2048, k=v=512. BIT-EXACT to 3 separate pwdqf16acc linears (fusing along N does not change the -// per-column K-accumulation order). Validated on Canary M4 Pro: correct (maxRel ~1e-3), scatter overhead -// 1.02x (free), concat win 1.63x on the QKV block. Boundaries hardcoded for Llama-3.2-1B GQA (32Q/8KV). + @group(0) @binding(0) var t_out_q: array; @group(0) @binding(1) var t_out_k: array; @group(0) @binding(2) var t_out_v: array; @@ -42,10 +35,12 @@ struct Params { _pad: u32, } @group(0) @binding(7) var params: Params; -const BM: u32 = 64u; const BN: u32 = 64u; const BK: u32 = 16u; + +// BK64 QKV variant: group_size=64 keeps one scale valid for all eight packed words. +const BM: u32 = 64u; const BN: u32 = 64u; const BK: u32 = 64u; const N_Q: u32 = 2048u; const N_QK: u32 = 2560u; const N_KV: u32 = 512u; -var As: array; -var Bs: array; +var As: array; +var Bs: array; @compute @workgroup_size(16, 16) fn main(@builtin(workgroup_id) wid: vec3, @builtin(local_invocation_id) lid: vec3) { @@ -61,18 +56,31 @@ fn main(@builtin(workgroup_id) wid: vec3, } let ar = tid / 4u; let ac = (tid % 4u) * 4u; + var k0: u32 = 0u; loop { if (k0 >= params.K) { break; } let arow = row0 + ar; if (arow < params.M) { let base = arow * params.K + k0 + ac; - let av = t_input[base >> 2u]; - As[ar * BK + ac + 0u] = f16(av.x); As[ar * BK + ac + 1u] = f16(av.y); - As[ar * BK + ac + 2u] = f16(av.z); As[ar * BK + ac + 3u] = f16(av.w); + let av0 = t_input[(base + 0u) >> 2u]; + let av1 = t_input[(base + 16u) >> 2u]; + let av2 = t_input[(base + 32u) >> 2u]; + let av3 = t_input[(base + 48u) >> 2u]; + As[ar * BK + ac + 0u] = f16(av0.x); As[ar * BK + ac + 1u] = f16(av0.y); + As[ar * BK + ac + 2u] = f16(av0.z); As[ar * BK + ac + 3u] = f16(av0.w); + As[ar * BK + ac + 16u] = f16(av1.x); As[ar * BK + ac + 17u] = f16(av1.y); + As[ar * BK + ac + 18u] = f16(av1.z); As[ar * BK + ac + 19u] = f16(av1.w); + As[ar * BK + ac + 32u] = f16(av2.x); As[ar * BK + ac + 33u] = f16(av2.y); + As[ar * BK + ac + 34u] = f16(av2.z); As[ar * BK + ac + 35u] = f16(av2.w); + As[ar * BK + ac + 48u] = f16(av3.x); As[ar * BK + ac + 49u] = f16(av3.y); + As[ar * BK + ac + 50u] = f16(av3.z); As[ar * BK + ac + 51u] = f16(av3.w); } else { - As[ar * BK + ac + 0u] = 0.0h; As[ar * BK + ac + 1u] = 0.0h; - As[ar * BK + ac + 2u] = 0.0h; As[ar * BK + ac + 3u] = 0.0h; + for (var segment: u32 = 0u; segment < 4u; segment = segment + 1u) { + for (var ai: u32 = 0u; ai < 4u; ai = ai + 1u) { + As[ar * BK + ac + segment * 16u + ai] = 0.0h; + } + } } if (tid < BN) { let c = tid; @@ -81,10 +89,17 @@ fn main(@builtin(workgroup_id) wid: vec3, let scale_row = (k0 / params.group_size) * params.padded_N; let scale = f16(t_scales[scale_row + n]); let base_word = n * (params.K_packed >> 2u) + (k0 >> 3u); - let w0 = t_weight[base_word]; + let w0 = t_weight[base_word + 0u]; let w1 = t_weight[base_word + 1u]; + let w2 = t_weight[base_word + 2u]; + let w3 = t_weight[base_word + 3u]; + let w4 = t_weight[base_word + 4u]; + let w5 = t_weight[base_word + 5u]; + let w6 = t_weight[base_word + 6u]; + let w7 = t_weight[base_word + 7u]; + let words = array(w0, w1, w2, w3, w4, w5, w6, w7); for (var br: u32 = 0u; br < BK; br = br + 1u) { - let word = select(w1, w0, br < 8u); + let word = words[br >> 3u]; let nib = (word >> ((br & 7u) * 4u)) & 0x0Fu; Bs[br * BN + c] = f16(i32(nib) - 8) * scale; } @@ -108,7 +123,7 @@ fn main(@builtin(workgroup_id) wid: vec3, for (var m: u32 = 0u; m < 4u; m = m + 1u) { for (var n: u32 = 0u; n < 4u; n = n + 1u) { let r = row0 + lid.y * 4u + m; - let c = col0 + lid.x * 4u + n; // global fused column [0, 3072) + let c = col0 + lid.x * 4u + n; if (r < params.M && c < params.N) { var val = f32(acc[m][n]); if (params.has_bias != 0u) { val = val + t_bias[c]; } @@ -121,8 +136,8 @@ fn main(@builtin(workgroup_id) wid: vec3, } )"; -inline constexpr uint32_t kQ4gswLinearGemmQkvFusedWorkgroupSizeX = 16; -inline constexpr uint32_t kQ4gswLinearGemmQkvFusedWorkgroupSizeY = 16; -inline constexpr uint32_t kQ4gswLinearGemmQkvFusedWorkgroupSizeZ = 1; +inline constexpr uint32_t kQ4gswQkvBk64WorkgroupSizeX = 16; +inline constexpr uint32_t kQ4gswQkvBk64WorkgroupSizeY = 16; +inline constexpr uint32_t kQ4gswQkvBk64WorkgroupSizeZ = 1; } // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_steel_bk64_wgsl.h b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_steel_bk64_wgsl.h new file mode 100644 index 00000000000..e502bf5a5ff --- /dev/null +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_steel_bk64_wgsl.h @@ -0,0 +1,158 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from q4gsw_linear_gemm_steel.wgsl - DO NOT EDIT. +// wgsl-sha256: dc25d61a4fe98a3fcfe9c006e7488ad84e2e873e6f66f26ac6cf41fd400da148 +inline constexpr const char* kQ4gswSteelBk64WGSL = R"( +enable f16; +@group(0) @binding(0) var t_out: array>; +@group(0) @binding(1) var t_input: array>; +@group(0) @binding(2) var t_weight: array; +@group(0) @binding(3) var t_scales: array; +@group(0) @binding(4) var t_bias: array; + +struct Params { + M: u32, + N: u32, + K: u32, + K_packed: u32, + group_size: u32, + padded_N: u32, + has_bias: u32, + _pad: u32, +} +@group(0) @binding(5) var params: Params; + +// "steel" prefill GEMM (M>1): 64x64 tile, 256 threads; K%16==0 host-guarded. +// The "steel" name + register-tiled dequant-to-shared GEMM structure are +// inspired by MLX's steel GEMM kernels (github.com/ml-explore/mlx, +// mlx/backend/metal/kernels/steel). One template, four variants: +// DTYPE=float f32 storage/multiply, per-nibble weight staging. +// DTYPE=half f16 storage/multiply, per-nibble weight staging. +// PWDQ (half only) packed-word dequant: load each u32 weight word ONCE, +// unpack all 16 nibbles of a column + hoist the per-column scale to one read +// (the per-nibble path re-reads each word ~8x). Requires K%BK==0 (steel +// route guarantees it) and group_size%BK==0 (hoisted scale across the tile). +// ACC=half (PWDQ only) f16 accumulate with fma(), cast to f32 in the epilogue +// -- LOSSY, perplexity-gated, opt-in via a runtime spec. ACC=float is f32 +// accumulate -- BIT-EXACT to the per-nibble half kernel. +// BK=64 (PWDQ + ACC=half only) stages a full quantization group at once. +const BM: u32 = 64u; const BN: u32 = 64u; const BK: u32 = 64u; +var As: array; // BM*BK +var Bs: array; // BK*BN +// 16x16 = 256 threads, bound to the 64x64 tile + 4x4 reg tile (not a knob). +@compute @workgroup_size(16, 16) +fn main(@builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3) { + let nbN = (params.N + BN - 1u) / BN; + let bx = wid.x % nbN; // decode 2D tile id from 1D dispatch + let by = wid.x / nbN; + let row0 = by * BM; + let col0 = bx * BN; + let tid = lid.y * 16u + lid.x; + var acc: array, 4>; + for (var m: u32 = 0u; m < 4u; m = m + 1u) { + for (var n: u32 = 0u; n < 4u; n = n + 1u) { acc[m][n] = 0.0h; } + } + // A staging coords: 256 threads load 64x16 = 1024 f32 -> 4 rows each (4 contiguous K). + let ar = tid / 4u; // 0..63 (row in tile) + let ac = (tid % 4u) * 4u; // 0,4,8,12 (K offset, 4 contiguous) + + var k0: u32 = 0u; + loop { + if (k0 >= params.K) { break; } + // stage activations (edge-masked on M; K is a multiple of BK for our shapes) + let arow = row0 + ar; + if (arow < params.M) { + let base = arow * params.K + k0 + ac; + // vec4 coalesced load; base is 4-aligned on the steel route (K%16==0, ac/k0 multiples of 4). + let av0 = t_input[(base + 0u) >> 2u]; + let av1 = t_input[(base + 16u) >> 2u]; + let av2 = t_input[(base + 32u) >> 2u]; + let av3 = t_input[(base + 48u) >> 2u]; + As[ar * BK + ac + 0u] = f16(av0.x); As[ar * BK + ac + 1u] = f16(av0.y); + As[ar * BK + ac + 2u] = f16(av0.z); As[ar * BK + ac + 3u] = f16(av0.w); + As[ar * BK + ac + 16u] = f16(av1.x); As[ar * BK + ac + 17u] = f16(av1.y); + As[ar * BK + ac + 18u] = f16(av1.z); As[ar * BK + ac + 19u] = f16(av1.w); + As[ar * BK + ac + 32u] = f16(av2.x); As[ar * BK + ac + 33u] = f16(av2.y); + As[ar * BK + ac + 34u] = f16(av2.z); As[ar * BK + ac + 35u] = f16(av2.w); + As[ar * BK + ac + 48u] = f16(av3.x); As[ar * BK + ac + 49u] = f16(av3.y); + As[ar * BK + ac + 50u] = f16(av3.z); As[ar * BK + ac + 51u] = f16(av3.w); + } else { + for (var segment: u32 = 0u; segment < 4u; segment = segment + 1u) { + for (var ai: u32 = 0u; ai < 4u; ai = ai + 1u) { + As[ar * BK + ac + segment * 16u + ai] = 0.0h; + } + } + } + // Packed-word dequant: threads [0,BN) each stage one full BK-column of Bs. + if (tid < BN) { + let c = tid; // Bs column within this tile + let n = col0 + c; // global output column + if (n < params.N) { + // Scale is constant across the BK tile (group_size % BK == 0 for all real + // group sizes; K%BK==0 on the steel route), so hoist it to one read. + let scale_row = (k0 / params.group_size) * params.padded_N; + let scale = f16(t_scales[scale_row + n]); + // Column n's BK-nibble K-slice starts at this packed word. + // K_packed multiple of 8 => base_word stays inside column n's own region. + let base_word = n * (params.K_packed >> 2u) + (k0 >> 3u); + let words = array( + t_weight[base_word + 0u], t_weight[base_word + 1u], + t_weight[base_word + 2u], t_weight[base_word + 3u], + t_weight[base_word + 4u], t_weight[base_word + 5u], + t_weight[base_word + 6u], t_weight[base_word + 7u]); + for (var br: u32 = 0u; br < BK; br = br + 1u) { + let word = words[br >> 3u]; + let nib = (word >> ((br & 7u) * 4u)) & 0x0Fu; + Bs[br * BN + c] = f16(i32(nib) - 8) * scale; + } + } else { + for (var br: u32 = 0u; br < BK; br = br + 1u) { Bs[br * BN + c] = 0.0h; } + } + } + workgroupBarrier(); + for (var k: u32 = 0u; k < BK; k = k + 1u) { + var a: array; + var bvec: array; + for (var m: u32 = 0u; m < 4u; m = m + 1u) { a[m] = As[(lid.y * 4u + m) * BK + k]; } + for (var n: u32 = 0u; n < 4u; n = n + 1u) { bvec[n] = Bs[k * BN + lid.x * 4u + n]; } + for (var m: u32 = 0u; m < 4u; m = m + 1u) { + for (var n: u32 = 0u; n < 4u; n = n + 1u) { acc[m][n] = fma(a[m], bvec[n], acc[m][n]); } + } + } + workgroupBarrier(); + k0 = k0 + BK; + } + for (var m: u32 = 0u; m < 4u; m = m + 1u) { + let r = row0 + lid.y * 4u + m; + let c0 = col0 + lid.x * 4u; + if (r < params.M && c0 < params.N) { + var vv = vec4( + f32(acc[m][0]), f32(acc[m][1]), f32(acc[m][2]), f32(acc[m][3])); + if (params.has_bias != 0u) { + vv = vv + vec4( + t_bias[c0], t_bias[c0 + 1u], t_bias[c0 + 2u], t_bias[c0 + 3u]); + } + t_out[(r * params.N + c0) >> 2u] = vv; + } + } +} +)"; + +inline constexpr uint32_t kQ4gswSteelBk64WorkgroupSizeX = 16; +inline constexpr uint32_t kQ4gswSteelBk64WorkgroupSizeY = 16; +inline constexpr uint32_t kQ4gswSteelBk64WorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/rope/RotaryEmbedding.cpp b/backends/webgpu/runtime/ops/rope/RotaryEmbedding.cpp index 012875b21e7..29f00c5823f 100644 --- a/backends/webgpu/runtime/ops/rope/RotaryEmbedding.cpp +++ b/backends/webgpu/runtime/ops/rope/RotaryEmbedding.cpp @@ -7,23 +7,22 @@ */ #include +#include #include #include -#include -#include #include #include -#include -#include #include -#include namespace executorch::backends::webgpu { namespace { +constexpr const char* kRotaryShader = "rotary_embedding"; +constexpr const char* kRotaryHfShader = "rotary_embedding_hf"; + // Uniform layout matching the WGSL Params struct (16-byte aligned, 32 bytes). struct RotaryParams { uint32_t n_heads; @@ -37,184 +36,109 @@ struct RotaryParams { }; static_assert(sizeof(RotaryParams) == 32, "RotaryParams must be 32 bytes"); -// A rope dispatch: its param-uniform (rewritten on resize) and its index in the -// graph's dispatch list (so a resize hook can update the workgroup count). -struct RopeDispatch { - WGPUBuffer uniform; - size_t dispatch_index; +enum class RopeGridPolicy { + OneDimensional, + FoldedTwoDimensional, +}; + +struct RopeGridContext { + int tensor_id; + uint32_t workgroup_size; + RopeGridPolicy policy; + const char* op_name; }; -// Rotate one (x->out) with the shared shader; freqs shared between xq and xk. -RopeDispatch add_rope_dispatch( +WebGPUDispatchGrid pick_rope_grid( + const WebGPUGraph& graph, + const RopeGridContext& context) { + const uint32_t num_pairs = static_cast( + utils::numel_of(graph.cur_dims(context.tensor_id)) / 2u); + if (context.policy == RopeGridPolicy::OneDimensional) { + return { + utils::compute_1d_workgroup_count( + graph.device(), num_pairs, context.workgroup_size, context.op_name), + 1u}; + } + const utils::WgCount grid = utils::compute_2d_workgroup_count( + graph.device(), num_pairs, context.workgroup_size, context.op_name); + return {grid.x, grid.y}; +} + +void preflight_rope_grids( + const WebGPUGraph& graph, + const RopeGridContext& q_context, + const RopeGridContext& k_context) { + (void)pick_rope_grid(graph, q_context); + (void)pick_rope_grid(graph, k_context); +} + +template +WGPUBuffer add_rope_dispatch( WebGPUGraph& graph, - WGPUDevice device, - std::optional& shared_resources, - uint32_t wg_size, + const char* shader_name, + const char* kernel_name, const WebGPUTensor& x, const WebGPUTensor& out, const WebGPUTensor& freqs_cos, const WebGPUTensor& freqs_sin, - uint32_t n_heads, - uint32_t seq, - uint32_t head_dim, - uint32_t workgroup_count) { - const uint32_t half_dim = head_dim / 2u; - // out.dims == in.dims (asserted in impl), so this matches the caller's wgc. - const uint32_t num_pairs = - static_cast(utils::numel_of(out.dims) / 2u); - - RotaryParams params = {}; - params.n_heads = n_heads; - params.seq = seq; - params.head_dim = head_dim; - params.half_dim = half_dim; - params.num_pairs = num_pairs; - - WGPUBufferDescriptor uniform_desc = {}; - uniform_desc.size = sizeof(RotaryParams); - uniform_desc.usage = WGPUBufferUsage_Uniform | WGPUBufferUsage_CopyDst; - uniform_desc.mappedAtCreation = true; - WGPUBuffer uniform_buffer = wgpuDeviceCreateBuffer(device, &uniform_desc); - void* mapped = - wgpuBufferGetMappedRange(uniform_buffer, 0, sizeof(RotaryParams)); - std::memcpy(mapped, ¶ms, sizeof(RotaryParams)); - wgpuBufferUnmap(uniform_buffer); - graph.add_uniform_buffer_bytes(sizeof(RotaryParams)); - - WGPUConstantEntry wg_size_constant = {}; - wg_size_constant.key = {"wg_size", WGPU_STRLEN}; - wg_size_constant.value = static_cast(wg_size); - - const std::vector bindings = { - {0, WGPUBufferBindingType_Storage, out.buffer, out.nbytes}, - {1, WGPUBufferBindingType_ReadOnlyStorage, x.buffer, x.nbytes}, - {2, - WGPUBufferBindingType_ReadOnlyStorage, - freqs_cos.buffer, - freqs_cos.nbytes}, - {3, - WGPUBufferBindingType_ReadOnlyStorage, - freqs_sin.buffer, - freqs_sin.nbytes}, - {4, WGPUBufferBindingType_Uniform, uniform_buffer, sizeof(RotaryParams)}, - }; - utils::ComputePipelineBundle bundle = shared_resources.has_value() - ? utils::make_compute_pipeline( - device, *shared_resources, bindings, &wg_size_constant, 1) - : utils::make_compute_pipeline( - device, kRotaryEmbeddingWGSL, bindings, &wg_size_constant, 1); - - const size_t dispatch_index = graph.add_dispatch( - {bundle.pipeline, - bundle.bind_group, - workgroup_count, - "apply_rotary_emb"}); - if (!shared_resources.has_value()) { - shared_resources.emplace(std::move(bundle)); - } - - // Graph owns it so a resize hook can rewrite it; freed in the dtor. - graph.own_uniform_buffer(uniform_buffer); - return {uniform_buffer, dispatch_index}; + const Params& params, + int trigger_tensor_id, + const RopeGridContext& grid_context, + uint32_t wg_size) { + WGPUBuffer uniform_buffer = graph.create_params_buffer(params); + WebGPUComputeDispatchDescriptor descriptor; + descriptor.shader_name = shader_name; + descriptor.kernel_name = kernel_name; + descriptor.bindings = { + {out.buffer, 0u, out.nbytes}, + {x.buffer, 0u, x.nbytes}, + {freqs_cos.buffer, 0u, freqs_cos.nbytes}, + {freqs_sin.buffer, 0u, freqs_sin.nbytes}, + {uniform_buffer, 0u, sizeof(Params)}}; + descriptor.constants = {{"wg_size", static_cast(wg_size)}}; + graph.add_dynamic_compute_dispatch( + descriptor, trigger_tensor_id, pick_rope_grid, grid_context); + return uniform_buffer; } -// Resize hook body: recompute S/num_pairs + both dispatches; out follows xq/xk. -void resize_rope( - WebGPUGraph& g, - int xq_id, - int xk_id, - int xq_out_id, - int xk_out_id, - uint32_t n_heads_q, - uint32_t n_heads_k, - uint32_t head_dim, - uint32_t half_dim, - uint32_t wg_size, - size_t q_idx, - size_t k_idx, - WGPUBuffer q_ubuf, - WGPUBuffer k_ubuf) { - const auto& qd = g.cur_dims(xq_id); - const auto& kd = g.cur_dims(xk_id); - if (qd.size() < 3 || kd.size() < 3) { - throw std::runtime_error("apply_rotary_emb(resize): q/k rank must be >= 3"); - } - const uint32_t s = static_cast(qd[qd.size() - 3]); - const uint64_t qn = utils::numel_of(qd); - const uint64_t kn = utils::numel_of(kd); - // pk = pq (seq=s); require k's seq == s, not silently q's. - if (static_cast(kd[kd.size() - 3]) != s) { - throw std::runtime_error( - "apply_rotary_emb(resize): q and k seq lengths differ"); - } - // freqs stay max-allocated; shader indexes by position (S = prefix). - RotaryParams pq = {}; - pq.n_heads = n_heads_q; - pq.seq = s; - pq.head_dim = head_dim; - pq.half_dim = half_dim; - pq.num_pairs = static_cast(qn / 2u); - RotaryParams pk = pq; - pk.n_heads = n_heads_k; - pk.num_pairs = static_cast(kn / 2u); - wgpuQueueWriteBuffer(g.queue(), q_ubuf, 0, &pq, sizeof(pq)); - wgpuQueueWriteBuffer(g.queue(), k_ubuf, 0, &pk, sizeof(pk)); - g.dispatch_at(q_idx).workgroup_count_x = utils::compute_1d_workgroup_count( - g.device(), - static_cast(qn / 2u), - wg_size, - "apply_rotary_emb(resize)"); - g.dispatch_at(k_idx).workgroup_count_x = utils::compute_1d_workgroup_count( - g.device(), - static_cast(kn / 2u), - wg_size, - "apply_rotary_emb(resize)"); - g.set_cur_dims(xq_out_id, qd); - g.set_cur_dims(xk_out_id, kd); -} - -// args: [xq, xk, freqs_cos, freqs_sin, out_list(ValueList[xq_out, xk_out])]. -void apply_rotary_emb_impl(WebGPUGraph& graph, const std::vector& args) { - const int xq_id = args.at(0); - const int xk_id = args.at(1); - const int freqs_cos_id = args.at(2); - const int freqs_sin_id = args.at(3); - - const std::vector& out_list = graph.get_value_list(args.at(4)); - if (out_list.size() != 2) { - throw std::runtime_error( - "WebGPU apply_rotary_emb: expected an output ValueList of size 2"); - } - - WGPUDevice device = graph.device(); - - const auto& xq = graph.get_tensor(xq_id); - const auto& xk = graph.get_tensor(xk_id); - const auto& freqs_cos = graph.get_tensor(freqs_cos_id); - const auto& freqs_sin = graph.get_tensor(freqs_sin_id); - const auto& xq_out = graph.get_tensor(out_list[0]); - const auto& xk_out = graph.get_tensor(out_list[1]); +struct RotaryGeometry { + uint32_t head_dim; + uint32_t seq; + uint32_t n_heads_q; + uint32_t n_heads_k; + uint32_t half_dim; + uint64_t xq_numel; + uint64_t xk_numel; +}; - // Vulkan shape contract: xq/xk (B,S,n_heads,head_dim), freqs (S,head_dim/2). +RotaryGeometry validate_rope_inputs( + const WebGPUTensor& xq, + const WebGPUTensor& xk, + const WebGPUTensor& freqs_cos, + const WebGPUTensor& freqs_sin, + const WebGPUTensor& xq_out, + const WebGPUTensor& xk_out) { if (xq.dims.size() < 3 || xk.dims.size() < 3 || freqs_cos.dims.size() < 2) { throw std::runtime_error("WebGPU apply_rotary_emb: malformed dims"); } - const uint32_t head_dim = static_cast(xq.dims.back()); - const uint32_t seq = static_cast(xq.dims[xq.dims.size() - 3]); - const uint32_t n_heads_q = static_cast(xq.dims[xq.dims.size() - 2]); - const uint32_t n_heads_k = static_cast(xk.dims[xk.dims.size() - 2]); + RotaryGeometry geometry = {}; + geometry.head_dim = static_cast(xq.dims.back()); + geometry.seq = static_cast(xq.dims[xq.dims.size() - 3]); + geometry.n_heads_q = static_cast(xq.dims[xq.dims.size() - 2]); + geometry.n_heads_k = static_cast(xk.dims[xk.dims.size() - 2]); const uint32_t seq_k = static_cast(xk.dims[xk.dims.size() - 3]); - const uint32_t half_dim = static_cast(freqs_cos.dims.back()); + geometry.half_dim = static_cast(freqs_cos.dims.back()); - if (head_dim == 0 || head_dim % 2 != 0) { + if (geometry.head_dim == 0 || geometry.head_dim % 2 != 0) { throw std::runtime_error( "WebGPU apply_rotary_emb: head_dim must be a nonzero multiple of 2"); } - if (static_cast(xk.dims.back()) != head_dim || seq_k != seq) { + if (static_cast(xk.dims.back()) != geometry.head_dim || + seq_k != geometry.seq) { throw std::runtime_error( "WebGPU apply_rotary_emb: xq/xk head_dim and seq must match"); } - if (half_dim * 2u != head_dim) { + if (geometry.half_dim * 2u != geometry.head_dim) { throw std::runtime_error( "WebGPU apply_rotary_emb: head_dim != 2 * freqs_cos last dim"); } @@ -222,127 +146,163 @@ void apply_rotary_emb_impl(WebGPUGraph& graph, const std::vector& args) { throw std::runtime_error( "WebGPU apply_rotary_emb: freqs_cos and freqs_sin shapes differ"); } - if (xq.buffer == nullptr || xk.buffer == nullptr || freqs_cos.buffer == nullptr || freqs_sin.buffer == nullptr || xq_out.buffer == nullptr || xk_out.buffer == nullptr) { throw std::runtime_error("WebGPU apply_rotary_emb: null buffer binding"); } - // All tensors are fp32; output shapes equal their inputs. - const uint64_t xq_numel = utils::numel_of(xq.dims); - const uint64_t xk_numel = utils::numel_of(xk.dims); + geometry.xq_numel = utils::numel_of(xq.dims); + geometry.xk_numel = utils::numel_of(xk.dims); const uint64_t freqs_numel = utils::numel_of(freqs_cos.dims); - if (freqs_numel != static_cast(seq) * half_dim || - xq.nbytes != xq_numel * sizeof(float) || - xk.nbytes != xk_numel * sizeof(float) || + if (freqs_numel != static_cast(geometry.seq) * geometry.half_dim || + xq.nbytes != geometry.xq_numel * sizeof(float) || + xk.nbytes != geometry.xk_numel * sizeof(float) || freqs_cos.nbytes != freqs_numel * sizeof(float) || freqs_sin.nbytes != freqs_numel * sizeof(float) || - xq_out.nbytes != xq_numel * sizeof(float) || - xk_out.nbytes != xk_numel * sizeof(float)) { + xq_out.nbytes != geometry.xq_numel * sizeof(float) || + xk_out.nbytes != geometry.xk_numel * sizeof(float)) { throw std::runtime_error( "WebGPU apply_rotary_emb: dtype/byte-size mismatch (all fp32) or " "freqs shape != [seq, head_dim/2]"); } + if (geometry.xq_numel > UINT32_MAX || geometry.xk_numel > UINT32_MAX) { + throw std::runtime_error( + "WebGPU apply_rotary_emb: element index exceeds uint32 range"); + } + return geometry; +} - if (xq_numel / 2u > UINT32_MAX || xk_numel / 2u > UINT32_MAX) { +struct RotaryResizeContext { + int xq_id; + int xk_id; + int xq_out_id; + int xk_out_id; + uint32_t n_heads_q; + uint32_t n_heads_k; + uint32_t head_dim; + uint32_t half_dim; + WGPUBuffer q_uniform; + WGPUBuffer k_uniform; +}; + +// Resize hook body: update parameters and outputs; the graph owns grid refresh. +void resize_rope(WebGPUGraph& graph, const RotaryResizeContext& context) { + const auto& q_dims = graph.cur_dims(context.xq_id); + const auto& k_dims = graph.cur_dims(context.xk_id); + if (q_dims.size() < 3 || k_dims.size() < 3) { + throw std::runtime_error("apply_rotary_emb(resize): q/k rank must be >= 3"); + } + const uint32_t seq = static_cast(q_dims[q_dims.size() - 3]); + const uint64_t q_numel = utils::numel_of(q_dims); + const uint64_t k_numel = utils::numel_of(k_dims); + // pk = pq (seq=s); require k's seq == s, not silently q's. + if (static_cast(k_dims[k_dims.size() - 3]) != seq) { throw std::runtime_error( - "WebGPU apply_rotary_emb: pair count exceeds uint32 dispatch range"); + "apply_rotary_emb(resize): q and k seq lengths differ"); } + // freqs stay max-allocated; shader indexes by position (S = prefix). + RotaryParams q_params = {}; + q_params.n_heads = context.n_heads_q; + q_params.seq = seq; + q_params.head_dim = context.head_dim; + q_params.half_dim = context.half_dim; + q_params.num_pairs = static_cast(q_numel / 2u); + RotaryParams k_params = q_params; + k_params.n_heads = context.n_heads_k; + k_params.num_pairs = static_cast(k_numel / 2u); + wgpuQueueWriteBuffer( + graph.queue(), context.q_uniform, 0, &q_params, sizeof(q_params)); + wgpuQueueWriteBuffer( + graph.queue(), context.k_uniform, 0, &k_params, sizeof(k_params)); + graph.set_cur_dims(context.xq_out_id, q_dims); + graph.set_cur_dims(context.xk_out_id, k_dims); +} - const uint32_t wg_size = - utils::clamp_workgroup_size(device, kRotaryEmbeddingWorkgroupSizeX); - // Validate both dispatches before any GPU-object alloc (no leak on throw). - const uint32_t xq_wgc = utils::compute_1d_workgroup_count( - device, - static_cast(xq_numel / 2u), - wg_size, - "apply_rotary_emb"); - const uint32_t xk_wgc = utils::compute_1d_workgroup_count( - device, - static_cast(xk_numel / 2u), - wg_size, - "apply_rotary_emb"); +// args: [xq, xk, freqs_cos, freqs_sin, out_list(ValueList[xq_out, xk_out])]. +void apply_rotary_emb_impl(WebGPUGraph& graph, const std::vector& args) { + const int xq_id = args.at(0); + const int xk_id = args.at(1); + const int freqs_cos_id = args.at(2); + const int freqs_sin_id = args.at(3); + + const std::vector& out_list = graph.get_value_list(args.at(4)); + if (out_list.size() != 2) { + throw std::runtime_error( + "WebGPU apply_rotary_emb: expected an output ValueList of size 2"); + } - std::optional shared_resources; - RopeDispatch q_disp = add_rope_dispatch( + const auto& xq = graph.get_tensor(xq_id); + const auto& xk = graph.get_tensor(xk_id); + const auto& freqs_cos = graph.get_tensor(freqs_cos_id); + const auto& freqs_sin = graph.get_tensor(freqs_sin_id); + const auto& xq_out = graph.get_tensor(out_list[0]); + const auto& xk_out = graph.get_tensor(out_list[1]); + + const RotaryGeometry geometry = + validate_rope_inputs(xq, xk, freqs_cos, freqs_sin, xq_out, xk_out); + + const uint32_t wg_size = utils::clamp_workgroup_size( + graph.device(), get_webgpu_shader_info(kRotaryShader).workgroup_size_x); + const RopeGridContext q_grid = { + xq_id, wg_size, RopeGridPolicy::OneDimensional, "apply_rotary_emb"}; + const RopeGridContext k_grid = { + xk_id, wg_size, RopeGridPolicy::OneDimensional, "apply_rotary_emb"}; + preflight_rope_grids(graph, q_grid, k_grid); + + RotaryParams q_params = {}; + q_params.n_heads = geometry.n_heads_q; + q_params.seq = geometry.seq; + q_params.head_dim = geometry.head_dim; + q_params.half_dim = geometry.half_dim; + q_params.num_pairs = static_cast(geometry.xq_numel / 2u); + RotaryParams k_params = q_params; + k_params.n_heads = geometry.n_heads_k; + k_params.num_pairs = static_cast(geometry.xk_numel / 2u); + const WGPUBuffer q_uniform = add_rope_dispatch( graph, - device, - shared_resources, - wg_size, + kRotaryShader, + "apply_rotary_emb", xq, xq_out, freqs_cos, freqs_sin, - n_heads_q, - seq, - head_dim, - xq_wgc); - RopeDispatch k_disp = add_rope_dispatch( + q_params, + xq_id, + q_grid, + wg_size); + const WGPUBuffer k_uniform = add_rope_dispatch( graph, - device, - shared_resources, - wg_size, + kRotaryShader, + "apply_rotary_emb", xk, xk_out, freqs_cos, freqs_sin, - n_heads_k, - seq, - head_dim, - xk_wgc); - WGPUBuffer q_ubuf = q_disp.uniform; - WGPUBuffer k_ubuf = k_disp.uniform; - const size_t q_idx = q_disp.dispatch_index; - const size_t k_idx = k_disp.dispatch_index; - - // Dynamic shapes: recompute S/num_pairs + both dispatches; out follows xq/xk. - const int xq_out_id = out_list[0]; - const int xk_out_id = out_list[1]; + k_params, + xk_id, + k_grid, + wg_size); + // Register on both xq and xk so the recompute fires whichever is marked dirty // (q and k co-resize on S; resize_rope is idempotent, so a double-fire when // both are dirty is harmless). - auto rope_hook = [xq_id, - xk_id, - xq_out_id, - xk_out_id, - n_heads_q, - n_heads_k, - head_dim, - half_dim, - wg_size, - q_idx, - k_idx, - q_ubuf, - k_ubuf](WebGPUGraph& g) { - resize_rope( - g, - xq_id, - xk_id, - xq_out_id, - xk_out_id, - n_heads_q, - n_heads_k, - head_dim, - half_dim, - wg_size, - q_idx, - k_idx, - q_ubuf, - k_ubuf); - }; - graph.add_tensor_resize_hook(xq_id, rope_hook); - graph.add_tensor_resize_hook(xk_id, rope_hook); + const RotaryResizeContext resize_context = { + xq_id, + xk_id, + out_list[0], + out_list[1], + geometry.n_heads_q, + geometry.n_heads_k, + geometry.head_dim, + geometry.half_dim, + q_uniform, + k_uniform}; + graph.add_tensor_resize_hook(xq_id, resize_rope, resize_context); + graph.add_tensor_resize_hook(xk_id, resize_rope, resize_context); } -// HuggingFace rotate-half RoPE (Qwen3 etc.). Structural sibling of the -// interleaved path above (same one-thread-per-pair scalar dispatch, wg_size, -// and resize hook); differs only in element pairing (i with i+half_dim vs -// even/odd), a full [max_seq, rotary_dim] freqs table, and a start_pos offset. -// Mirrors Vulkan's et_vk.apply_rotary_emb_hf -// (backends/vulkan/runtime/graph/ops/impl/RotaryEmbedding.cpp:211). - -// Uniform layout matching the HF WGSL Params struct (32 bytes). +// Mirrors Vulkan's full-dimension HuggingFace rotate-half RoPE. struct RotaryHfParams { uint32_t n_heads; uint32_t seq; @@ -355,187 +315,65 @@ struct RotaryHfParams { }; static_assert(sizeof(RotaryHfParams) == 32, "RotaryHfParams must be 32 bytes"); -RopeDispatch add_rope_hf_dispatch( - WebGPUGraph& graph, - uint32_t wg_size, - const WebGPUTensor& x, - const WebGPUTensor& out, - const WebGPUTensor& freqs_cos, - const WebGPUTensor& freqs_sin, - uint32_t n_heads, - uint32_t seq, - uint32_t head_dim, - uint32_t half_dim, - uint32_t rotary_dim, - uint32_t start_pos, - uint32_t workgroup_count) { - const uint32_t num_pairs = - static_cast(utils::numel_of(out.dims) / 2u); - - RotaryHfParams params = {}; - params.n_heads = n_heads; - params.seq = seq; - params.head_dim = head_dim; - params.half_dim = half_dim; - params.num_pairs = num_pairs; - params.rotary_dim = rotary_dim; - params.start_pos = start_pos; - - WGPUBuffer uniform_buffer = - utils::make_uniform(graph.device(), ¶ms, sizeof(RotaryHfParams)); - graph.add_uniform_buffer_bytes(sizeof(RotaryHfParams)); - - WGPUConstantEntry wg_size_constant = {}; - wg_size_constant.key = {"wg_size", WGPU_STRLEN}; - wg_size_constant.value = static_cast(wg_size); - - utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( - graph.device(), - kRotaryEmbeddingHfWGSL, - { - {0, WGPUBufferBindingType_Storage, out.buffer, out.nbytes}, - {1, WGPUBufferBindingType_ReadOnlyStorage, x.buffer, x.nbytes}, - {2, - WGPUBufferBindingType_ReadOnlyStorage, - freqs_cos.buffer, - freqs_cos.nbytes}, - {3, - WGPUBufferBindingType_ReadOnlyStorage, - freqs_sin.buffer, - freqs_sin.nbytes}, - {4, - WGPUBufferBindingType_Uniform, - uniform_buffer, - sizeof(RotaryHfParams)}, - }, - &wg_size_constant, - 1); - - const size_t dispatch_index = graph.add_dispatch( - {bundle.pipeline, - bundle.bind_group, - workgroup_count, - "apply_rotary_emb_hf"}); - - graph.own_uniform_buffer(uniform_buffer); - return {uniform_buffer, dispatch_index}; -} - -// Resize hook body: recompute S/num_pairs + (dynamic) start_pos for both -// dispatches; out follows xq/xk. Fires on xq/xk seq resize and, when start_pos -// is a runtime SymInt (KV-cache decode), on each start_pos change; idempotent. -void resize_rope_hf( - WebGPUGraph& g, - int xq_id, - int xk_id, - int xq_out_id, - int xk_out_id, - int start_pos_id, - bool dynamic_pos, - uint32_t baked_start_pos, - uint32_t n_heads_q, - uint32_t n_heads_k, - uint32_t head_dim, - uint32_t half_dim, - uint32_t rotary_dim, - uint32_t wg_size, - size_t q_idx, - size_t k_idx, - WGPUBuffer q_ubuf, - WGPUBuffer k_ubuf) { - const auto& qd = g.cur_dims(xq_id); - const auto& kd = g.cur_dims(xk_id); - if (qd.size() < 3 || kd.size() < 3) { - throw std::runtime_error( - "apply_rotary_emb_hf(resize): q/k rank must be >= 3"); - } - const uint32_t s = static_cast(qd[qd.size() - 3]); - const uint64_t qn = utils::numel_of(qd); - const uint64_t kn = utils::numel_of(kd); - if (static_cast(kd[kd.size() - 3]) != s) { - throw std::runtime_error( - "apply_rotary_emb_hf(resize): q and k seq lengths differ"); - } - uint32_t start_pos = baked_start_pos; - if (dynamic_pos) { - const int32_t pos = g.read_symint(start_pos_id); - if (pos < 0) { - throw std::runtime_error( - "apply_rotary_emb_hf(resize): start_pos must be non-negative"); - } - start_pos = static_cast(pos); - } - RotaryHfParams pq = {}; - pq.n_heads = n_heads_q; - pq.seq = s; - pq.head_dim = head_dim; - pq.half_dim = half_dim; - pq.num_pairs = static_cast(qn / 2u); - pq.rotary_dim = rotary_dim; - pq.start_pos = start_pos; - RotaryHfParams pk = pq; - pk.n_heads = n_heads_k; - pk.num_pairs = static_cast(kn / 2u); - wgpuQueueWriteBuffer(g.queue(), q_ubuf, 0, &pq, sizeof(pq)); - wgpuQueueWriteBuffer(g.queue(), k_ubuf, 0, &pk, sizeof(pk)); - g.dispatch_at(q_idx).workgroup_count_x = utils::compute_1d_workgroup_count( - g.device(), - static_cast(qn / 2u), - wg_size, - "apply_rotary_emb_hf(resize)"); - g.dispatch_at(k_idx).workgroup_count_x = utils::compute_1d_workgroup_count( - g.device(), - static_cast(kn / 2u), - wg_size, - "apply_rotary_emb_hf(resize)"); - g.set_cur_dims(xq_out_id, qd); - g.set_cur_dims(xk_out_id, kd); -} - -// Validated HF-rope shape, derived from the input tensors. -struct RopeHfShape { +struct RotaryHfGeometry { uint32_t head_dim; uint32_t seq; uint32_t n_heads_q; uint32_t n_heads_k; + uint32_t max_seq; uint32_t rotary_dim; uint32_t half_dim; uint64_t xq_numel; uint64_t xk_numel; }; -// Derive and validate the HF-rope input shapes; throws on any malformed input. -RopeHfShape validate_rope_hf_inputs( - const WebGPUTensor& xq, +RotaryHfGeometry validate_rope_hf_inputs( + const WebGPUTensor& x, const WebGPUTensor& xk, const WebGPUTensor& freqs_cos, const WebGPUTensor& freqs_sin, - const WebGPUTensor& xq_out, + const WebGPUTensor& x_out, const WebGPUTensor& xk_out) { - // Shape contract: xq/xk (B,S,n_heads,head_dim), freqs (max_seq, rotary_dim). - if (xq.dims.size() < 3 || xk.dims.size() < 3 || freqs_cos.dims.size() < 2) { + if (x.dims.size() < 3 || xk.dims.size() != x.dims.size() || + freqs_cos.dims.size() != 2) { throw std::runtime_error("WebGPU apply_rotary_emb_hf: malformed dims"); } - const uint32_t head_dim = static_cast(xq.dims.back()); - const uint32_t seq = static_cast(xq.dims[xq.dims.size() - 3]); - const uint32_t n_heads_q = static_cast(xq.dims[xq.dims.size() - 2]); - const uint32_t n_heads_k = static_cast(xk.dims[xk.dims.size() - 2]); - const uint32_t seq_k = static_cast(xk.dims[xk.dims.size() - 3]); - const uint32_t max_seq = - static_cast(freqs_cos.dims[freqs_cos.dims.size() - 2]); - const uint32_t rotary_dim = static_cast(freqs_cos.dims.back()); - - if (head_dim == 0 || head_dim % 2 != 0) { + if (x_out.dims != x.dims || xk_out.dims != xk.dims) { + throw std::runtime_error( + "WebGPU apply_rotary_emb_hf: output shapes must match q/k inputs"); + } + for (size_t i = 0; i + 3 < x.dims.size(); i++) { + if (x.dims[i] != xk.dims[i]) { + throw std::runtime_error( + "WebGPU apply_rotary_emb_hf: q/k batch dimensions differ"); + } + } + const auto positive_u32 = [](int64_t value, const char* label) { + if (value <= 0 || static_cast(value) > UINT32_MAX) { + throw std::runtime_error( + std::string("WebGPU apply_rotary_emb_hf: invalid ") + label); + } + return static_cast(value); + }; + RotaryHfGeometry geometry = {}; + geometry.head_dim = positive_u32(x.dims.back(), "head_dim"); + geometry.seq = positive_u32(x.dims[x.dims.size() - 3], "sequence length"); + geometry.n_heads_q = + positive_u32(x.dims[x.dims.size() - 2], "query head count"); + geometry.n_heads_k = + positive_u32(xk.dims[xk.dims.size() - 2], "key head count"); + geometry.max_seq = positive_u32(freqs_cos.dims[0], "frequency row count"); + geometry.rotary_dim = positive_u32(freqs_cos.dims[1], "rotary_dim"); + if (geometry.head_dim % 2 != 0) { throw std::runtime_error( "WebGPU apply_rotary_emb_hf: head_dim must be a nonzero multiple of 2"); } - if (static_cast(xk.dims.back()) != head_dim || seq_k != seq) { + if (xk.dims.back() != static_cast(geometry.head_dim) || + xk.dims[xk.dims.size() - 3] != static_cast(geometry.seq)) { throw std::runtime_error( "WebGPU apply_rotary_emb_hf: xq/xk head_dim and seq must match"); } - // Full rotary only (rotary_dim == head_dim); partial-rotary passthrough is a - // documented follow-up (Qwen3 uses full RoPE). Throw rather than mis-rotate. - if (rotary_dim != head_dim) { + if (geometry.rotary_dim != geometry.head_dim) { throw std::runtime_error( "WebGPU apply_rotary_emb_hf: partial rotary (rotary_dim != head_dim) " "not supported"); @@ -544,51 +382,134 @@ RopeHfShape validate_rope_hf_inputs( throw std::runtime_error( "WebGPU apply_rotary_emb_hf: freqs_cos and freqs_sin shapes differ"); } - if (max_seq < seq) { + if (geometry.max_seq < geometry.seq) { throw std::runtime_error("WebGPU apply_rotary_emb_hf: freqs max_seq < seq"); } - - if (xq.buffer == nullptr || xk.buffer == nullptr || + if (x.buffer == nullptr || xk.buffer == nullptr || freqs_cos.buffer == nullptr || freqs_sin.buffer == nullptr || - xq_out.buffer == nullptr || xk_out.buffer == nullptr) { + x_out.buffer == nullptr || xk_out.buffer == nullptr) { throw std::runtime_error("WebGPU apply_rotary_emb_hf: null buffer binding"); } + const WebGPUTensor* tensors[] = { + &x, &xk, &freqs_cos, &freqs_sin, &x_out, &xk_out}; + for (const WebGPUTensor* tensor : tensors) { + if (tensor->is_int || tensor->elem_size != sizeof(float)) { + throw std::runtime_error( + "WebGPU apply_rotary_emb_hf: all tensors must be fp32"); + } + } - // All tensors are fp32; output shapes equal their inputs. - const uint64_t xq_numel = utils::numel_of(xq.dims); - const uint64_t xk_numel = utils::numel_of(xk.dims); + geometry.half_dim = geometry.rotary_dim / 2u; + geometry.xq_numel = utils::numel_of(x.dims); + geometry.xk_numel = utils::numel_of(xk.dims); const uint64_t freqs_numel = utils::numel_of(freqs_cos.dims); - if (freqs_numel != static_cast(max_seq) * rotary_dim || - xq.nbytes != xq_numel * sizeof(float) || - xk.nbytes != xk_numel * sizeof(float) || + if (freqs_numel != + static_cast(geometry.max_seq) * geometry.rotary_dim || + x.nbytes != geometry.xq_numel * sizeof(float) || + xk.nbytes != geometry.xk_numel * sizeof(float) || freqs_cos.nbytes != freqs_numel * sizeof(float) || freqs_sin.nbytes != freqs_numel * sizeof(float) || - xq_out.nbytes != xq_numel * sizeof(float) || - xk_out.nbytes != xk_numel * sizeof(float)) { + x_out.nbytes != geometry.xq_numel * sizeof(float) || + xk_out.nbytes != geometry.xk_numel * sizeof(float)) { throw std::runtime_error( "WebGPU apply_rotary_emb_hf: dtype/byte-size mismatch (all fp32) or " "freqs shape != [max_seq, rotary_dim]"); } + if (geometry.xq_numel == 0 || geometry.xk_numel == 0 || + geometry.xq_numel > UINT32_MAX || geometry.xk_numel > UINT32_MAX) { + throw std::runtime_error( + "WebGPU apply_rotary_emb_hf: element index exceeds uint32 range"); + } + return geometry; +} + +struct RotaryHfResizeContext { + int xq_id; + int xk_id; + int xq_out_id; + int xk_out_id; + int start_pos_id; + bool dynamic_pos; + uint32_t baked_start_pos; + uint32_t n_heads_q; + uint32_t n_heads_k; + uint32_t head_dim; + uint32_t half_dim; + uint32_t rotary_dim; + uint32_t max_seq; + WGPUBuffer q_uniform; + WGPUBuffer k_uniform; +}; - if (xq_numel / 2u > UINT32_MAX || xk_numel / 2u > UINT32_MAX) { +void resize_rope_hf(WebGPUGraph& graph, const RotaryHfResizeContext& context) { + const auto& q_dims = graph.cur_dims(context.xq_id); + const auto& k_dims = graph.cur_dims(context.xk_id); + if (q_dims.size() < 3 || k_dims.size() != q_dims.size()) { + throw std::runtime_error( + "apply_rotary_emb_hf(resize): q/k rank must be >= 3"); + } + const int64_t seq_value = q_dims[q_dims.size() - 3]; + if (seq_value <= 0 || static_cast(seq_value) > UINT32_MAX) { + throw std::runtime_error( + "apply_rotary_emb_hf(resize): invalid sequence length"); + } + const uint32_t seq = static_cast(seq_value); + if (k_dims[k_dims.size() - 3] != seq_value) { + throw std::runtime_error( + "apply_rotary_emb_hf(resize): q and k seq lengths differ"); + } + if (q_dims.back() != static_cast(context.head_dim) || + k_dims.back() != static_cast(context.head_dim) || + q_dims[q_dims.size() - 2] != static_cast(context.n_heads_q) || + k_dims[k_dims.size() - 2] != static_cast(context.n_heads_k)) { throw std::runtime_error( - "WebGPU apply_rotary_emb_hf: pair count exceeds uint32 dispatch range"); - } - - return { - head_dim, - seq, - n_heads_q, - n_heads_k, - rotary_dim, - rotary_dim / 2u, - xq_numel, - xk_numel}; + "apply_rotary_emb_hf(resize): q/k head geometry changed"); + } + for (size_t i = 0; i + 3 < q_dims.size(); i++) { + if (q_dims[i] != k_dims[i]) { + throw std::runtime_error( + "apply_rotary_emb_hf(resize): q/k batch dimensions differ"); + } + } + const uint64_t q_numel = utils::numel_of(q_dims); + const uint64_t k_numel = utils::numel_of(k_dims); + if (q_numel == 0 || k_numel == 0 || q_numel > UINT32_MAX || + k_numel > UINT32_MAX) { + throw std::runtime_error( + "apply_rotary_emb_hf(resize): element index exceeds uint32 range"); + } + uint32_t start_pos = context.baked_start_pos; + if (context.dynamic_pos) { + const int64_t pos = graph.read_symint(context.start_pos_id); + if (pos < 0 || static_cast(pos) > UINT32_MAX) { + throw std::runtime_error( + "apply_rotary_emb_hf(resize): start_pos must be non-negative"); + } + start_pos = static_cast(pos); + } + if (static_cast(start_pos) + seq > context.max_seq) { + throw std::runtime_error( + "apply_rotary_emb_hf(resize): start_pos + seq exceeds freqs max_seq"); + } + RotaryHfParams q_params = {}; + q_params.n_heads = context.n_heads_q; + q_params.seq = seq; + q_params.head_dim = context.head_dim; + q_params.half_dim = context.half_dim; + q_params.num_pairs = static_cast(q_numel / 2u); + q_params.rotary_dim = context.rotary_dim; + q_params.start_pos = start_pos; + RotaryHfParams k_params = q_params; + k_params.n_heads = context.n_heads_k; + k_params.num_pairs = static_cast(k_numel / 2u); + wgpuQueueWriteBuffer( + graph.queue(), context.q_uniform, 0, &q_params, sizeof(q_params)); + wgpuQueueWriteBuffer( + graph.queue(), context.k_uniform, 0, &k_params, sizeof(k_params)); + graph.set_cur_dims(context.xq_out_id, q_dims); + graph.set_cur_dims(context.xk_out_id, k_dims); } -// args: [xq, xk, freqs_cos, freqs_sin, start_pos, out_list(ValueList[xq_out, -// xk_out])]. freqs is the FULL [max_seq, rotary_dim] table (start_pos offsets -// into it), unlike the pre-sliced interleaved freqs. void apply_rotary_emb_hf_impl( WebGPUGraph& graph, const std::vector& args) { @@ -604,8 +525,6 @@ void apply_rotary_emb_hf_impl( "WebGPU apply_rotary_emb_hf: expected an output ValueList of size 2"); } - WGPUDevice device = graph.device(); - const auto& xq = graph.get_tensor(xq_id); const auto& xk = graph.get_tensor(xk_id); const auto& freqs_cos = graph.get_tensor(freqs_cos_id); @@ -613,24 +532,15 @@ void apply_rotary_emb_hf_impl( const auto& xq_out = graph.get_tensor(out_list[0]); const auto& xk_out = graph.get_tensor(out_list[1]); - const RopeHfShape shp = + const RotaryHfGeometry geometry = validate_rope_hf_inputs(xq, xk, freqs_cos, freqs_sin, xq_out, xk_out); - const uint32_t head_dim = shp.head_dim; - const uint32_t seq = shp.seq; - const uint32_t n_heads_q = shp.n_heads_q; - const uint32_t n_heads_k = shp.n_heads_k; - const uint32_t rotary_dim = shp.rotary_dim; - const uint32_t half_dim = shp.half_dim; - const uint64_t xq_numel = shp.xq_numel; - const uint64_t xk_numel = shp.xk_numel; - - // start_pos: build-time Int (baked) OR runtime SymInt (dynamic decode); - // mirrors sdpa's input_pos handling. + + // Decode uses a SymInt position; static graphs use an Int. int64_t start_pos = 0; const auto start_pos_type = graph.get_value_type(start_pos_id); const bool dynamic_pos = start_pos_type == WebGPUGraph::ValueType::SymInt; if (dynamic_pos) { - start_pos = graph.read_symint(start_pos_id); // build placeholder (e.g. 0) + start_pos = graph.read_symint(start_pos_id); } else if (start_pos_type == WebGPUGraph::ValueType::Int) { start_pos = graph.get_int(start_pos_id); } else { @@ -641,99 +551,82 @@ void apply_rotary_emb_hf_impl( throw std::runtime_error( "WebGPU apply_rotary_emb_hf: start_pos must be non-negative"); } + if (static_cast(start_pos) + geometry.seq > geometry.max_seq) { + throw std::runtime_error( + "WebGPU apply_rotary_emb_hf: start_pos + seq exceeds freqs max_seq"); + } - const uint32_t wg_size = - utils::clamp_workgroup_size(device, kRotaryEmbeddingHfWorkgroupSizeX); - // Validate both dispatches before any GPU-object alloc (no leak on throw). - const uint32_t xq_wgc = utils::compute_1d_workgroup_count( - device, - static_cast(xq_numel / 2u), + const uint32_t wg_size = utils::clamp_workgroup_size( + graph.device(), get_webgpu_shader_info(kRotaryHfShader).workgroup_size_x); + const RopeGridContext q_grid = { + xq_id, wg_size, - "apply_rotary_emb_hf"); - const uint32_t xk_wgc = utils::compute_1d_workgroup_count( - device, - static_cast(xk_numel / 2u), + RopeGridPolicy::FoldedTwoDimensional, + "apply_rotary_emb_hf"}; + const RopeGridContext k_grid = { + xk_id, wg_size, - "apply_rotary_emb_hf"); - - RopeDispatch q_disp = add_rope_hf_dispatch( + RopeGridPolicy::FoldedTwoDimensional, + "apply_rotary_emb_hf"}; + preflight_rope_grids(graph, q_grid, k_grid); + + RotaryHfParams q_params = {}; + q_params.n_heads = geometry.n_heads_q; + q_params.seq = geometry.seq; + q_params.head_dim = geometry.head_dim; + q_params.half_dim = geometry.half_dim; + q_params.num_pairs = static_cast(geometry.xq_numel / 2u); + q_params.rotary_dim = geometry.rotary_dim; + q_params.start_pos = static_cast(start_pos); + RotaryHfParams k_params = q_params; + k_params.n_heads = geometry.n_heads_k; + k_params.num_pairs = static_cast(geometry.xk_numel / 2u); + + const WGPUBuffer q_uniform = add_rope_dispatch( graph, - wg_size, + kRotaryHfShader, + "apply_rotary_emb_hf", xq, xq_out, freqs_cos, freqs_sin, - n_heads_q, - seq, - head_dim, - half_dim, - rotary_dim, - static_cast(start_pos), - xq_wgc); - RopeDispatch k_disp = add_rope_hf_dispatch( + q_params, + xq_id, + q_grid, + wg_size); + const WGPUBuffer k_uniform = add_rope_dispatch( graph, - wg_size, + kRotaryHfShader, + "apply_rotary_emb_hf", xk, xk_out, freqs_cos, freqs_sin, - n_heads_k, - seq, - head_dim, - half_dim, - rotary_dim, + k_params, + xk_id, + k_grid, + wg_size); + + const RotaryHfResizeContext resize_context = { + xq_id, + xk_id, + out_list[0], + out_list[1], + start_pos_id, + dynamic_pos, static_cast(start_pos), - xk_wgc); - WGPUBuffer q_ubuf = q_disp.uniform; - WGPUBuffer k_ubuf = k_disp.uniform; - const size_t q_idx = q_disp.dispatch_index; - const size_t k_idx = k_disp.dispatch_index; - - const int xq_out_id = out_list[0]; - const int xk_out_id = out_list[1]; - const uint32_t baked_start_pos = static_cast(start_pos); - auto rope_hook = [xq_id, - xk_id, - xq_out_id, - xk_out_id, - start_pos_id, - dynamic_pos, - baked_start_pos, - n_heads_q, - n_heads_k, - head_dim, - half_dim, - rotary_dim, - wg_size, - q_idx, - k_idx, - q_ubuf, - k_ubuf](WebGPUGraph& g) { - resize_rope_hf( - g, - xq_id, - xk_id, - xq_out_id, - xk_out_id, - start_pos_id, - dynamic_pos, - baked_start_pos, - n_heads_q, - n_heads_k, - head_dim, - half_dim, - rotary_dim, - wg_size, - q_idx, - k_idx, - q_ubuf, - k_ubuf); - }; - graph.add_tensor_resize_hook(xq_id, rope_hook); - graph.add_tensor_resize_hook(xk_id, rope_hook); - // Dynamic decode: re-fire when the runtime start_pos SymInt changes. + geometry.n_heads_q, + geometry.n_heads_k, + geometry.head_dim, + geometry.half_dim, + geometry.rotary_dim, + geometry.max_seq, + q_uniform, + k_uniform}; + graph.add_tensor_resize_hook(xq_id, resize_rope_hf, resize_context); + graph.add_tensor_resize_hook(xk_id, resize_rope_hf, resize_context); if (dynamic_pos) { - graph.add_resize_hook(start_pos_id, rope_hook); + graph.add_resize_hook(start_pos_id, resize_rope_hf, resize_context); } } diff --git a/backends/webgpu/runtime/ops/rope/rotary_embedding_hf.wgsl b/backends/webgpu/runtime/ops/rope/rotary_embedding_hf.wgsl index 14a6853afa3..859b24ca6fc 100644 --- a/backends/webgpu/runtime/ops/rope/rotary_embedding_hf.wgsl +++ b/backends/webgpu/runtime/ops/rope/rotary_embedding_hf.wgsl @@ -18,11 +18,13 @@ struct Params { override wg_size: u32 = 64u; // One thread per (i, i+half_dim) pair; HuggingFace rotate-half RoPE, shared -// xq/xk shader. freqs is the FULL [max_seq, rotary_dim] table (duplicated -// halves) indexed at row (start_pos + s); only the first-half column is read. +// xq/xk shader. freqs is the FULL [max_seq, rotary_dim] table indexed at row +// (start_pos + s); each output half uses its corresponding frequency column. @compute @workgroup_size(wg_size, 1, 1) -fn main(@builtin(global_invocation_id) gid: vec3) { - let pair = gid.x; +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let pair = gid.x + gid.y * (num_workgroups.x * wg_size); if (pair >= params.num_pairs) { return; } @@ -38,12 +40,16 @@ fn main(@builtin(global_invocation_id) gid: vec3) { ((b * params.seq + s) * params.n_heads + head) * params.head_dim; let a_idx = head_base + pair_i; let b_idx = head_base + pair_i + half_dim; - let freqs_idx = (s + params.start_pos) * params.rotary_dim + pair_i; + let freqs_base = (s + params.start_pos) * params.rotary_dim; + let freqs_a_idx = freqs_base + pair_i; + let freqs_b_idx = freqs_a_idx + half_dim; - let c = t_freqs_cos[freqs_idx]; - let si = t_freqs_sin[freqs_idx]; + let c_a = t_freqs_cos[freqs_a_idx]; + let si_a = t_freqs_sin[freqs_a_idx]; + let c_b = t_freqs_cos[freqs_b_idx]; + let si_b = t_freqs_sin[freqs_b_idx]; let x_a = t_in[a_idx]; let x_b = t_in[b_idx]; - t_out[a_idx] = x_a * c - x_b * si; - t_out[b_idx] = x_b * c + x_a * si; + t_out[a_idx] = x_a * c_a - x_b * si_a; + t_out[b_idx] = x_b * c_b + x_a * si_b; } diff --git a/backends/webgpu/runtime/ops/rope/rotary_embedding_hf_wgsl.h b/backends/webgpu/runtime/ops/rope/rotary_embedding_hf_wgsl.h index 191ec710e66..21242fb480a 100644 --- a/backends/webgpu/runtime/ops/rope/rotary_embedding_hf_wgsl.h +++ b/backends/webgpu/runtime/ops/rope/rotary_embedding_hf_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from rotary_embedding_hf.wgsl - DO NOT EDIT. -// wgsl-sha256: 5ba8d45925f00f12af17bf3092a1af9513a9e501c5c35e6b0d48cfb3dac7b5d6 +// wgsl-sha256: 4f081ed4c8165f021cbb722d379e437f30b8dfb08bf03bfcbaa406ed7799c7b6 inline constexpr const char* kRotaryEmbeddingHfWGSL = R"( @group(0) @binding(0) var t_out: array; @group(0) @binding(1) var t_in: array; @@ -35,11 +35,13 @@ struct Params { override wg_size: u32 = 64u; // One thread per (i, i+half_dim) pair; HuggingFace rotate-half RoPE, shared -// xq/xk shader. freqs is the FULL [max_seq, rotary_dim] table (duplicated -// halves) indexed at row (start_pos + s); only the first-half column is read. +// xq/xk shader. freqs is the FULL [max_seq, rotary_dim] table indexed at row +// (start_pos + s); each output half uses its corresponding frequency column. @compute @workgroup_size(wg_size, 1, 1) -fn main(@builtin(global_invocation_id) gid: vec3) { - let pair = gid.x; +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let pair = gid.x + gid.y * (num_workgroups.x * wg_size); if (pair >= params.num_pairs) { return; } @@ -55,14 +57,18 @@ fn main(@builtin(global_invocation_id) gid: vec3) { ((b * params.seq + s) * params.n_heads + head) * params.head_dim; let a_idx = head_base + pair_i; let b_idx = head_base + pair_i + half_dim; - let freqs_idx = (s + params.start_pos) * params.rotary_dim + pair_i; + let freqs_base = (s + params.start_pos) * params.rotary_dim; + let freqs_a_idx = freqs_base + pair_i; + let freqs_b_idx = freqs_a_idx + half_dim; - let c = t_freqs_cos[freqs_idx]; - let si = t_freqs_sin[freqs_idx]; + let c_a = t_freqs_cos[freqs_a_idx]; + let si_a = t_freqs_sin[freqs_a_idx]; + let c_b = t_freqs_cos[freqs_b_idx]; + let si_b = t_freqs_sin[freqs_b_idx]; let x_a = t_in[a_idx]; let x_b = t_in[b_idx]; - t_out[a_idx] = x_a * c - x_b * si; - t_out[b_idx] = x_b * c + x_a * si; + t_out[a_idx] = x_a * c_a - x_b * si_a; + t_out[b_idx] = x_b * c_b + x_a * si_b; } )"; diff --git a/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp b/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp index 96ec782c72e..bcb5deaa770 100644 --- a/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp +++ b/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp @@ -6,17 +6,12 @@ * LICENSE file in the root directory of this source tree. */ +#include #include +#include #include #include -#include -#include -#include -#include -#include #include -#include -#include #include @@ -34,6 +29,20 @@ namespace { constexpr int64_t kSdpaTileM = 4; constexpr int64_t kSdpaTileN = 4; +constexpr const char* kUpdateCacheShader = "update_cache"; +constexpr const char* kUpdateCacheHalfShader = "update_cache_half"; +constexpr const char* kAttnWeightsShader = "sdpa_compute_attn_weights"; +constexpr const char* kAttnWeightsHalfShader = "sdpa_compute_attn_weights_half"; +constexpr const char* kSoftmaxShader = "sdpa_softmax"; +constexpr const char* kComputeOutShader = "sdpa_compute_out"; +constexpr const char* kComputeOutHalfShader = "sdpa_compute_out_half"; +constexpr const char* kStreamingK16Shader = + "streaming_attention_k16_causal_bound"; +constexpr const char* kStreamingQwen3K16Shader = + "streaming_attention_qwen3_k16_causal_bound"; +constexpr const char* kStreamingQwen3Q32K16Shader = + "streaming_attention_qwen3_q32_k16_causal_bound"; + // Uniform param structs (all 16-byte aligned, matching the WGSL Params). struct UpdateCacheParams { uint32_t numel; @@ -75,6 +84,42 @@ struct ComputeOutParams { }; static_assert(sizeof(ComputeOutParams) == 32, "ComputeOutParams must be 32B"); +struct StreamingAttentionK16Params { + uint32_t S; + uint32_t context_len; + uint32_t input_pos; + uint32_t q_token_stride4; + uint32_t q_head_stride4; + uint32_t kv_token_stride4; + uint32_t kv_head_stride4; + uint32_t o_token_stride4; + uint32_t o_head_stride4; + uint32_t _pad0; + uint32_t _pad1; + uint32_t _pad2; +}; +static_assert( + sizeof(StreamingAttentionK16Params) == 48, + "StreamingAttentionK16Params must be 48B"); + +struct SdpaLiveState { + int64_t s; + int64_t pos; + int64_t context_len; + UpdateCacheParams update_cache; + AttnWeightsParams attn_weights; + SoftmaxParams softmax; + ComputeOutParams compute_out; + StreamingAttentionK16Params streaming_k16; + utils::WgCount update_cache_grid; + utils::WgCount qk_grid; + utils::WgCount softmax_grid; + utils::WgCount av_grid; + utils::WgCount streaming_k16_grid; + bool use_fd; + SdpaFdDecodeState fd; +}; + // Param-struct builder helpers — used in both initial build and resize hook. static UpdateCacheParams make_update_cache_params( uint64_t kv_numel, @@ -132,111 +177,176 @@ static ComputeOutParams make_compute_out_params( return p; } -// A buffer + its byte size, for binding. -struct BufferBinding { - WGPUBuffer buffer; - uint64_t size; -}; +static StreamingAttentionK16Params make_streaming_attention_k16_params( + int64_t S, + int64_t context_len, + int64_t input_pos, + int64_t Hq, + int64_t Hkv, + int64_t D) { + StreamingAttentionK16Params p = {}; + p.S = static_cast(S); + p.context_len = static_cast(context_len); + p.input_pos = static_cast(input_pos); + p.q_token_stride4 = static_cast(Hq * D / 4); + p.q_head_stride4 = static_cast(D / 4); + p.kv_token_stride4 = static_cast(Hkv * D / 4); + p.kv_head_stride4 = static_cast(D / 4); + p.o_token_stride4 = static_cast(Hq * D / 4); + p.o_head_stride4 = static_cast(D / 4); + return p; +} + +static bool streaming_attention_k16_device_supported(WGPUDevice device) { + WGPULimits limits = {}; + const WebGPUContext* context = get_default_webgpu_context(); + return context != nullptr && context->device == device && + context->shader_f16_supported && + wgpuDeviceGetLimits(device, &limits) == WGPUStatus_Success && + limits.maxComputeInvocationsPerWorkgroup >= 128u && + limits.maxComputeWorkgroupSizeX >= 32u && + limits.maxComputeWorkgroupSizeY >= 4u && + limits.maxComputeWorkgroupStorageSize >= 14720u; +} + +constexpr uint32_t kLlamaK16QueryTile = 32u; +constexpr uint32_t kQwen3K16QueryTile = 16u; +constexpr uint32_t kQwen3Q32K16QueryTile = 32u; +constexpr uint32_t kQwen3Q16K16StorageBytes = 512u * 4u * sizeof(float) + + 512u * 4u * sizeof(uint16_t) + 128u * 2u * sizeof(float) + + 3u * 16u * sizeof(float); +constexpr uint32_t kQwen3K16StorageBytes = kQwen3Q16K16StorageBytes; +// Mirrors the Q32 shader's workgroup arrays (t_q_tile vec4x1024, t_kv_tile +// vec4x512, t_scores vec2x256, t_m/t_d/t_alpha f32x32) so the +// device-support gate stays tied to the declared storage, not a literal. +constexpr uint32_t kQwen3Q32K16StorageBytes = 1024u * 4u * sizeof(float) + + 512u * 4u * sizeof(uint16_t) + 256u * 2u * sizeof(float) + + 3u * 32u * sizeof(float); + +constexpr bool streaming_attention_k16_workgroup_count_fits( + int64_t S, + int64_t Hkv, + int64_t g, + uint32_t query_tile, + uint32_t max_workgroups) { + if (S <= 0 || Hkv <= 0 || g <= 0 || query_tile == 0u || + max_workgroups == 0u) { + return false; + } + if (static_cast(S) > UINT64_MAX / static_cast(g)) { + return false; + } + const uint64_t logical_rows = + static_cast(S) * static_cast(g); + if (logical_rows > UINT64_MAX - (query_tile - 1u)) { + return false; + } + const uint64_t groups_per_kv = (logical_rows + query_tile - 1u) / query_tile; + if (groups_per_kv > UINT64_MAX / static_cast(Hkv)) { + return false; + } + const uint64_t workgroups = groups_per_kv * static_cast(Hkv); + return workgroups > 0u && workgroups <= UINT32_MAX && + workgroups <= max_workgroups; +} + +static_assert( + streaming_attention_k16_workgroup_count_fits(65528, 8, 2, 16, 65535)); +static_assert( + !streaming_attention_k16_workgroup_count_fits(65529, 8, 2, 16, 65535)); + +bool qwen3_q16_k16_device_supported(WGPUDevice device) { + WGPULimits limits = {}; + const WebGPUContext* context = get_default_webgpu_context(); + return context != nullptr && context->device == device && + context->shader_f16_supported && + wgpuDeviceGetLimits(device, &limits) == WGPUStatus_Success && + limits.maxComputeWorkgroupSizeX >= 16u && + limits.maxComputeWorkgroupSizeY >= 8u && + limits.maxComputeInvocationsPerWorkgroup >= 128u && + limits.maxComputeWorkgroupStorageSize >= kQwen3K16StorageBytes && + limits.maxStorageBuffersPerShaderStage >= 4u; +} -// Build one dispatch (pipeline + bind group) and record it on the graph. -void build_dispatch( +bool qwen3_q32_k16_device_supported(WGPUDevice device) { + WGPULimits limits = {}; + const WebGPUContext* context = get_default_webgpu_context(); + return context != nullptr && context->device == device && + context->shader_f16_supported && + wgpuDeviceGetLimits(device, &limits) == WGPUStatus_Success && + limits.maxComputeWorkgroupSizeX >= 32u && + limits.maxComputeWorkgroupSizeY >= 8u && + limits.maxComputeInvocationsPerWorkgroup >= 256u && + limits.maxComputeWorkgroupStorageSize >= kQwen3Q32K16StorageBytes && + limits.maxStorageBuffersPerShaderStage >= 4u; +} + +static utils::WgCount streaming_attention_k16_grid( + WGPUDevice device, + int64_t S, + int64_t Hkv, + int64_t g, + uint32_t query_tile) { + const uint64_t groups_per_kv = + (static_cast(S) * static_cast(g) + query_tile - 1u) / + query_tile; + const uint64_t workgroups = static_cast(Hkv) * groups_per_kv; + if (workgroups == 0u || workgroups > UINT32_MAX) { + throw std::runtime_error("WebGPU sdpa: K16 workgroup count exceeds uint32"); + } + if (workgroups > utils::queried_max_workgroups(device)) { + throw std::runtime_error( + "WebGPU sdpa: K16 workgroup count exceeds the 1D dispatch limit"); + } + return {static_cast(workgroups), 1u}; +} + +size_t add_sdpa_compute_dispatch( WebGPUGraph& graph, - const char* wgsl_source, - const BufferBinding* storage_bindings, - uint32_t n_storage, // includes the rw output at index 0 + const char* shader_name, + std::vector bindings, WGPUBuffer uniform_buffer, uint64_t uniform_size, - uint32_t workgroup_count_x, - uint32_t workgroup_count_y, + utils::WgCount grid, uint32_t wg_size, - bool retain_uniform = false, const char* kernel_name = "") { - WGPUDevice device = graph.device(); - - // Bind group layout: storage entries then the uniform. - constexpr uint32_t kMaxEntries = 8; - if (n_storage + 1 > kMaxEntries) { - throw std::runtime_error("WebGPU sdpa: n_storage exceeds kMaxEntries"); - } - const uint32_t uniform_binding = n_storage; - std::vector bindings; - bindings.reserve(n_storage + 1u); - for (uint32_t i = 0; i < n_storage; i++) { - bindings.push_back( - {i, - (i == 0) ? WGPUBufferBindingType_Storage - : WGPUBufferBindingType_ReadOnlyStorage, - storage_bindings[i].buffer, - storage_bindings[i].size}); - } - bindings.push_back( - {uniform_binding, - WGPUBufferBindingType_Uniform, - uniform_buffer, - uniform_size}); - - // All callers pass an override wg_size; a 0 would keep the shader default. - WGPUConstantEntry wg_size_constant = {}; - wg_size_constant.key = {"wg_size", WGPU_STRLEN}; - wg_size_constant.value = static_cast(wg_size); - - utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( - device, - wgsl_source, - bindings, - wg_size != 0 ? &wg_size_constant : nullptr, - wg_size != 0 ? 1u : 0u); - - graph.add_dispatch( - {bundle.pipeline, - bundle.bind_group, - workgroup_count_x, - kernel_name, - workgroup_count_y}); - - if (retain_uniform) { - // Graph owns it so a resize hook can rewrite it; freed in the dtor. - graph.own_uniform_buffer(uniform_buffer); - } else { - // Drop our ref; the bind group keeps the uniform alive. - wgpuBufferRelease(uniform_buffer); + bindings.push_back({uniform_buffer, 0u, uniform_size}); + WebGPUComputeDispatchDescriptor descriptor; + descriptor.shader_name = shader_name; + descriptor.kernel_name = kernel_name; + descriptor.bindings = std::move(bindings); + if (wg_size != 0) { + descriptor.constants = {{"wg_size", static_cast(wg_size)}}; } + descriptor.grid = {grid.x, grid.y}; + return graph.add_compute_dispatch(descriptor); } // Dispatch one update_cache (K or V); returns the retained uniform buffer. static WGPUBuffer record_update_cache_dispatch( WebGPUGraph& graph, - WGPUDevice device, const WebGPUTensor& cache, const WebGPUTensor& src, uint64_t kv_numel, uint32_t kv_dst_offset, uint64_t cache_numel, uint32_t uc_wg, - bool retain_uniform, const char* label) { const uint32_t wgc = utils::compute_1d_workgroup_count( - device, static_cast(kv_numel), uc_wg, label); - UpdateCacheParams uc = + graph.device(), static_cast(kv_numel), uc_wg, label); + const UpdateCacheParams uc = make_update_cache_params(kv_numel, kv_dst_offset, cache_numel); - WGPUBuffer ubuf = graph.make_uniform_buffer(&uc, sizeof(uc)); - BufferBinding bindings[2] = { - {cache.buffer, cache.nbytes}, {src.buffer, src.nbytes}}; - const char* uc_src = kUpdateCacheWGSL; - if (graph.kv_f16()) { - uc_src = kUpdateCacheHalfWGSL; - } - build_dispatch( + WGPUBuffer ubuf = graph.create_params_buffer(uc); + const std::vector bindings = { + {cache.buffer, 0u, cache.nbytes}, {src.buffer, 0u, src.nbytes}}; + add_sdpa_compute_dispatch( graph, - uc_src, + graph.kv_f16() ? kUpdateCacheHalfShader : kUpdateCacheShader, bindings, - 2, ubuf, sizeof(uc), - wgc, - 1, + {wgc, 1u}, uc_wg, - retain_uniform, "update_cache"); return ubuf; } @@ -280,10 +390,13 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { const size_t cn = k_cache.dims.size(); const int64_t Cmax = k_cache.dims[cn - 3]; - // Validate B == 1 (leading dims must all be 1). - for (size_t i = 0; i + 3 < qn; i++) { - if (q.dims[i] != 1) { - throw std::runtime_error("WebGPU sdpa: only batch size 1 is supported"); + // Validate B == 1 for every tensor (leading dims must all be 1). Rank-3 + // tensors are the equivalent squeezed-batch representation. + for (const WebGPUTensor* tensor : {&q, &k, &v, &k_cache, &v_cache, &out}) { + for (size_t i = 0; i + 3 < tensor->dims.size(); i++) { + if (tensor->dims[i] != 1) { + throw std::runtime_error("WebGPU sdpa: only batch size 1 is supported"); + } } } if (S <= 0 || Hq <= 0 || D <= 0 || Hkv <= 0 || Cmax <= 0) { @@ -319,8 +432,16 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { if (k_cache.dims != v_cache.dims) { throw std::runtime_error("WebGPU sdpa: k_cache and v_cache shape mismatch"); } + if (k_cache.dims[cn - 2] != Hkv) { + throw std::runtime_error( + "WebGPU sdpa: cache num_heads must match projected k/v"); + } + if (out.dims != q.dims) { + throw std::runtime_error("WebGPU sdpa: output shape must match q"); + } - // fp32-only: validate byte counts against fp32 element counts. + // q/k/v/out are serialized fp32. KV caches are fp32 by default and use + // dedicated fp16 storage only when the graph-level option is active. auto numel = [](const WebGPUTensor& t) { uint64_t n = 1; for (int64_t d : t.dims) { @@ -328,11 +449,23 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { } return n; }; - if (q.nbytes != numel(q) * sizeof(float) || - k.nbytes != numel(k) * sizeof(float) || - v.nbytes != numel(v) * sizeof(float) || - out.nbytes != numel(out) * sizeof(float)) { - throw std::runtime_error("WebGPU sdpa: fp32-only (byte-size mismatch)"); + auto is_fp32 = [&numel](const WebGPUTensor& t) { + return !t.is_int && t.elem_size == sizeof(float) && + t.nbytes == numel(t) * sizeof(float); + }; + if (!is_fp32(q) || !is_fp32(k) || !is_fp32(v) || !is_fp32(out)) { + throw std::runtime_error("WebGPU sdpa: q/k/v/output must be fp32"); + } + const size_t cache_elem_size = + graph.kv_f16() ? sizeof(uint16_t) : sizeof(float); + auto cache_storage_is_valid = [&numel, + cache_elem_size](const WebGPUTensor& t) { + return !t.is_int && t.elem_size == cache_elem_size && + t.nbytes == numel(t) * cache_elem_size; + }; + if (!cache_storage_is_valid(k_cache) || !cache_storage_is_valid(v_cache)) { + throw std::runtime_error( + "WebGPU sdpa: cache dtype does not match the selected storage mode"); } // input_pos: build-time Int (baked) OR runtime SymInt (dynamic decode). @@ -383,279 +516,439 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { throw std::runtime_error("WebGPU sdpa: only is_causal=true is supported"); } - // KV cache written in place; only attn_weights/softmax need scratch. - const uint64_t aw_floats = static_cast(Hq) * - static_cast(S) * static_cast(context_len); - // Dynamic input_pos: size+bind scratch for Cmax (no realloc; covers any ctx). - const uint64_t aw_cap_floats = static_cast(Hq) * - static_cast(S) * - static_cast(dynamic_pos ? Cmax : context_len); - const uint64_t aw_bytes = aw_cap_floats * sizeof(float); - - // Dynamic input_pos: the resize hook rewrites these per step. - WGPUBuffer uc_k_buf = nullptr, uc_v_buf = nullptr, qk_buf = nullptr, - softmax_buf = nullptr, av_buf = nullptr; - size_t qk_idx = 0, uc_k_idx = 0, uc_v_idx = 0, softmax_idx = 0, av_idx = 0; - const WGPUDevice device = graph.device(); - const uint32_t uc_wg = - utils::clamp_workgroup_size(device, kUpdateCacheWorkgroupSizeX); + const WGPUBuffer k16_buffers[] = { + q.buffer, k.buffer, v.buffer, k_cache.buffer, v_cache.buffer, out.buffer}; + bool k16_buffers_distinct = true; + for (size_t i = 0; i < 6; i++) { + for (size_t j = i + 1; j < 6; j++) { + k16_buffers_distinct = + k16_buffers_distinct && k16_buffers[i] != k16_buffers[j]; + } + } + // The specialized shaders bake the standard Qwen3 scale, so eligibility must + // be exact. A nearby explicit scale has different operator semantics and must + // use the general path. + const float qwen3_expected_scale = 1.0f / std::sqrt(128.0f); + const bool qwen3_k16_geometry = Hq == 16 && Hkv == 8 && g == 2 && D == 128 && + scale == qwen3_expected_scale && out.dims == q.dims; + // Q16 is the default route for exact Qwen3 geometry; Q32 is an explicit + // autotuning candidate requested via the sdpa_query_tile RuntimeSpec. Support + // is evaluated per-tile so an unsupported Q32 request falls back to the Q16 + // streaming route instead of dropping to the materialized path. + const uint32_t device_max_workgroups = utils::queried_max_workgroups(device); + const bool qwen3_q16_supported = + qwen3_k16_geometry && graph.kv_f16() && + qwen3_q16_k16_device_supported(device) && + streaming_attention_k16_workgroup_count_fits( + S, Hkv, g, kQwen3K16QueryTile, device_max_workgroups); + const bool qwen3_q32_requested = qwen3_k16_geometry && + graph.sdpa_query_tile() == static_cast(kQwen3Q32K16QueryTile); + const bool qwen3_q32_supported = + qwen3_q32_requested && graph.kv_f16() && + qwen3_q32_k16_device_supported(device) && + streaming_attention_k16_workgroup_count_fits( + S, Hkv, g, kQwen3Q32K16QueryTile, device_max_workgroups); + const bool qwen3_q32_selected = qwen3_q32_supported; + const bool qwen3_k16_selected = qwen3_q32_selected || qwen3_q16_supported; + const uint32_t qwen3_query_tile = + qwen3_q32_selected ? kQwen3Q32K16QueryTile : kQwen3K16QueryTile; + const bool llama_k16_eligible = + graph.kv_f16() && Hq == 32 && Hkv == 8 && g == 4 && D == 64 && + scale == 0.125f && out.dims == q.dims && + streaming_attention_k16_device_supported(device) && + streaming_attention_k16_workgroup_count_fits( + S, Hkv, g, kLlamaK16QueryTile, device_max_workgroups); + const bool k16_eligible = + k16_buffers_distinct && (llama_k16_eligible || qwen3_k16_selected); + const uint32_t k16_query_tile = + qwen3_k16_selected ? qwen3_query_tile : kLlamaK16QueryTile; + const char* k16_shader = qwen3_q32_selected ? kStreamingQwen3Q32K16Shader + : qwen3_k16_selected ? kStreamingQwen3K16Shader + : kStreamingK16Shader; + const char* k16_label = qwen3_q32_selected + ? "sdpa_streaming_attention_qwen3_q32_k16_causal_bound" + : qwen3_k16_selected ? "sdpa_streaming_attention_qwen3_k16_causal_bound" + : "sdpa_streaming_attention_k16_causal_bound"; + const uint32_t uc_wg = utils::clamp_workgroup_size( + device, get_webgpu_shader_info(kUpdateCacheShader).workgroup_size_x); const uint32_t qk_wg = utils::clamp_workgroup_size( - device, kSdpaComputeAttnWeightsWorkgroupSizeX); - const uint32_t av_wg = - utils::clamp_workgroup_size(device, kSdpaComputeOutWorkgroupSizeX); - - // Dispatches 1-2: write new K/V into the caches (reuses update_cache). - const uint64_t kv_numel = static_cast(S) * - static_cast(Hkv) * static_cast(D); - const uint32_t kv_dst_offset = static_cast( - static_cast(input_pos) * static_cast(Hkv) * - static_cast(D)); - uc_k_buf = record_update_cache_dispatch( + device, get_webgpu_shader_info(kAttnWeightsShader).workgroup_size_x); + const uint32_t av_wg = utils::clamp_workgroup_size( + device, get_webgpu_shader_info(kComputeOutShader).workgroup_size_x); + const uint32_t sm_wg = utils::clamp_workgroup_size_pow2( + device, get_webgpu_shader_info(kSoftmaxShader).workgroup_size_x); + const bool fd_eligible = D <= kSdpaFdMaxHeadDim; + const int64_t pos_const = input_pos; + + auto compute_live_state = [q_id, + k_id, + v_id, + out_id, + qn, + kn, + S, + dynamic_pos, + input_pos_id, + pos_const, + Hq, + Hkv, + D, + Cmax, + g, + scale, + uc_wg, + qk_wg, + av_wg, + fd_eligible, + k16_eligible, + k16_query_tile](WebGPUGraph& gr) { + SdpaLiveState state = {}; + const auto& q_live_dims = gr.cur_dims(q_id); + state.s = q_live_dims[qn - 3]; + state.pos = dynamic_pos ? static_cast(gr.read_symint(input_pos_id)) + : pos_const; + if (state.s <= 0 || state.pos < 0 || state.s > S) { + throw std::runtime_error("WebGPU sdpa: invalid live S or input_pos"); + } + if (gr.cur_dims(k_id)[kn - 3] != state.s || + gr.cur_dims(v_id)[gr.cur_dims(v_id).size() - 3] != state.s) { + throw std::runtime_error("WebGPU sdpa: live q/k/v seq_len mismatch"); + } + const auto& out_max_dims = gr.get_tensor(out_id).dims; + if (out_max_dims.size() != q_live_dims.size()) { + throw std::runtime_error("WebGPU sdpa: output rank must match q"); + } + for (size_t i = 0; i < q_live_dims.size(); i++) { + if (q_live_dims[i] <= 0 || q_live_dims[i] > out_max_dims[i]) { + throw std::runtime_error( + "WebGPU sdpa: live output shape exceeds allocation"); + } + } + state.context_len = state.s + state.pos; + if (state.context_len <= 0 || state.context_len > Cmax || + state.s > UINT32_MAX || state.pos > UINT32_MAX || + state.context_len > UINT32_MAX) { + throw std::runtime_error( + "WebGPU sdpa: live dimensions exceed cache or uint32 capacity"); + } + + const uint64_t kv_numel = static_cast(state.s) * + static_cast(Hkv) * static_cast(D); + const uint64_t kv_offset = static_cast(state.pos) * + static_cast(Hkv) * static_cast(D); + const uint64_t cache_numel = static_cast(Cmax) * + static_cast(Hkv) * static_cast(D); + if (kv_numel > UINT32_MAX || kv_offset > UINT32_MAX || + cache_numel > UINT32_MAX) { + throw std::runtime_error("WebGPU sdpa: live workload exceeds uint32"); + } + + state.update_cache = make_update_cache_params( + kv_numel, static_cast(kv_offset), cache_numel); + state.attn_weights = make_attn_weights_params( + state.s, Hq, Hkv, D, state.context_len, state.pos, g, scale); + state.softmax = make_softmax_params(Hq, state.s, state.context_len); + state.compute_out = + make_compute_out_params(state.s, Hq, Hkv, D, state.context_len, g); + if (k16_eligible) { + state.streaming_k16 = make_streaming_attention_k16_params( + state.s, state.context_len, state.pos, Hq, Hkv, D); + state.streaming_k16_grid = streaming_attention_k16_grid( + gr.device(), state.s, Hkv, g, k16_query_tile); + } + state.update_cache_grid = { + utils::compute_1d_workgroup_count( + gr.device(), static_cast(kv_numel), uc_wg, "uc(resize)"), + 1u}; + if (!k16_eligible) { + const uint64_t aw_floats = static_cast(Hq) * + static_cast(state.s) * + static_cast(state.context_len); + const uint64_t qk_tiles = static_cast(Hq) * + static_cast(utils::div_up(state.s, kSdpaTileM)) * + static_cast(utils::div_up(state.context_len, kSdpaTileN)); + const uint64_t softmax_rows = + static_cast(Hq) * static_cast(state.s); + const uint64_t av_tiles = static_cast(Hq) * + static_cast(utils::div_up(state.s, kSdpaTileM)) * + static_cast(utils::div_up(D, kSdpaTileN)); + if (aw_floats > UINT32_MAX || qk_tiles > UINT32_MAX || + softmax_rows > UINT32_MAX || av_tiles > UINT32_MAX) { + throw std::runtime_error( + "WebGPU sdpa: materialized workload exceeds uint32"); + } + state.qk_grid = utils::compute_2d_workgroup_count( + gr.device(), static_cast(qk_tiles), qk_wg, "QK(resize)"); + state.softmax_grid = utils::compute_2d_workgroup_count( + gr.device(), + static_cast(softmax_rows), + 1, + "softmax(resize)"); + state.av_grid = utils::compute_2d_workgroup_count( + gr.device(), static_cast(av_tiles), av_wg, "AV(resize)"); + } + state.use_fd = fd_eligible && state.s == 1; + // make_sdpa_fd_decode_state requires D % 4 == 0; the op-level guard above + // ("head_dim (D) must be a multiple of 4") rejects any other D before this + // lambda runs, so eager construction here can never throw on it. + if (fd_eligible) { + state.fd = make_sdpa_fd_decode_state( + gr.device(), Hq, Hkv, D, state.context_len, g, scale); + } + return state; + }; + + const SdpaLiveState initial_state = compute_live_state(graph); + const uint64_t aw_cap_floats = k16_eligible + ? 0u + : static_cast(Hq) * static_cast(S) * + static_cast(dynamic_pos ? Cmax : context_len); + const uint64_t aw_bytes = aw_cap_floats * sizeof(float); + + WGPUBuffer uc_k_buf = record_update_cache_dispatch( graph, - device, k_cache, k, - kv_numel, - kv_dst_offset, - numel(k_cache), + initial_state.update_cache.numel, + initial_state.update_cache.dst_offset, + initial_state.update_cache.cache_numel, uc_wg, - true, "update_cache(K)"); - uc_v_buf = record_update_cache_dispatch( + WGPUBuffer uc_v_buf = record_update_cache_dispatch( graph, - device, v_cache, v, - kv_numel, - kv_dst_offset, - numel(v_cache), + initial_state.update_cache.numel, + initial_state.update_cache.dst_offset, + initial_state.update_cache.cache_numel, uc_wg, - true, "update_cache(V)"); - uc_k_idx = graph.num_dispatches() - 2; - uc_v_idx = graph.num_dispatches() - 1; - - // FlashDecoding decode (S==1, static pos). Shapes FD can't handle (head dim - // > kSdpaFdMaxHeadDim) fall through to the materialized path below. - if (S == 1 && !dynamic_pos && D <= kSdpaFdMaxHeadDim) { - sdpa_fd_decode_dispatch( - graph, q, k_cache, v_cache, out, Hq, Hkv, D, context_len, g, scale); - return; - } - - // QK/softmax scratch — allocated only on the non-FD path (Hq*S*Cmax prefill). - WGPUBuffer attn_weights = graph.acquire_scratch(aw_bytes); - WebGPUGraph::ScopedScratch attn_weights_guard(&graph, attn_weights); - WGPUBuffer attn_weights_softmax = graph.acquire_scratch(aw_bytes); - WebGPUGraph::ScopedScratch attn_weights_softmax_guard( - &graph, attn_weights_softmax); - - // --- Dispatch 3: QK -> attn_weights. One thread per TM x TN tile. - { - if (aw_floats > UINT32_MAX) { - throw std::runtime_error( - "WebGPU sdpa: Hq*S*context_len exceeds uint32 max"); - } - const int64_t qk_tiles = Hq * utils::div_up(S, kSdpaTileM) * - utils::div_up(context_len, kSdpaTileN); - const utils::WgCount wgc = utils::compute_2d_workgroup_count( - device, static_cast(qk_tiles), qk_wg, "QK"); - AttnWeightsParams p = make_attn_weights_params( - S, Hq, Hkv, D, context_len, input_pos, g, scale); - WGPUBuffer ubuf = graph.make_uniform_buffer(&p, sizeof(p)); - BufferBinding bindings[3] = { - {attn_weights, aw_bytes}, - {q.buffer, q.nbytes}, - {k_cache.buffer, k_cache.nbytes}}; - const char* qk_src = kSdpaComputeAttnWeightsWGSL; - if (graph.kv_f16()) { - qk_src = kSdpaComputeAttnWeightsHalfWGSL; - } - build_dispatch( + const size_t uc_k_idx = graph.num_dispatches() - 2; + const size_t uc_v_idx = graph.num_dispatches() - 1; + const bool dynamic_sequence = graph.tensor_has_dynamic_dims(q_id) || + graph.tensor_has_dynamic_dims(k_id) || + graph.tensor_has_dynamic_dims(v_id); + const bool dual_route = utils::should_record_sdpa_dual_route( + fd_eligible, dynamic_sequence, dynamic_pos); + const bool record_k16 = k16_eligible && (dual_route || !initial_state.use_fd); + const bool record_materialized = + !k16_eligible && (dual_route || !initial_state.use_fd); + const bool record_fd = dual_route || initial_state.use_fd; + + WGPUBuffer qk_buf = nullptr; + WGPUBuffer softmax_buf = nullptr; + WGPUBuffer av_buf = nullptr; + size_t qk_idx = 0; + size_t softmax_idx = 0; + size_t av_idx = 0; + utils::DispatchRange materialized_range = {}; + if (record_materialized) { + WGPUBuffer attn_weights = graph.acquire_scratch(aw_bytes); + WebGPUGraph::ScopedScratch attn_weights_guard(&graph, attn_weights); + WGPUBuffer attn_weights_softmax = graph.acquire_scratch(aw_bytes); + WebGPUGraph::ScopedScratch attn_weights_softmax_guard( + &graph, attn_weights_softmax); + + materialized_range.begin = graph.num_dispatches(); + qk_buf = graph.create_params_buffer(initial_state.attn_weights); + const std::vector qk_bindings = { + {attn_weights, 0u, aw_bytes}, + {q.buffer, 0u, q.nbytes}, + {k_cache.buffer, 0u, k_cache.nbytes}}; + add_sdpa_compute_dispatch( graph, - qk_src, - bindings, - 3, - ubuf, - sizeof(p), - wgc.x, - wgc.y, + graph.kv_f16() ? kAttnWeightsHalfShader : kAttnWeightsShader, + qk_bindings, + qk_buf, + sizeof(AttnWeightsParams), + initial_state.qk_grid, qk_wg, - true, "sdpa_compute_attn_weights"); - qk_buf = ubuf; qk_idx = graph.num_dispatches() - 1; - } - // Dispatch 4: softmax, one workgroup per (h,s) row of width context_len. - { - // One workgroup per (h,s) row; wg_size 1 keeps the device dispatch check. - const utils::WgCount wgc = utils::compute_2d_workgroup_count( - device, static_cast(Hq * S), 1, "softmax"); - const uint32_t sm_wg = - utils::clamp_workgroup_size_pow2(device, kSdpaSoftmaxWorkgroupSizeX); - SoftmaxParams p = make_softmax_params(Hq, S, context_len); - WGPUBuffer ubuf = graph.make_uniform_buffer(&p, sizeof(p)); - BufferBinding bindings[2] = { - {attn_weights_softmax, aw_bytes}, {attn_weights, aw_bytes}}; - build_dispatch( + softmax_buf = graph.create_params_buffer(initial_state.softmax); + const std::vector softmax_bindings = { + {attn_weights_softmax, 0u, aw_bytes}, {attn_weights, 0u, aw_bytes}}; + add_sdpa_compute_dispatch( graph, - kSdpaSoftmaxWGSL, - bindings, - 2, - ubuf, - sizeof(p), - wgc.x, - wgc.y, + kSoftmaxShader, + softmax_bindings, + softmax_buf, + sizeof(SoftmaxParams), + initial_state.softmax_grid, sm_wg, - true, "sdpa_softmax"); - softmax_buf = ubuf; softmax_idx = graph.num_dispatches() - 1; - } - // --- Dispatch 5: AV -> out. One thread per TM x TN tile. - { - const int64_t av_tiles = - Hq * utils::div_up(S, kSdpaTileM) * utils::div_up(D, kSdpaTileN); - const utils::WgCount wgc = utils::compute_2d_workgroup_count( - device, static_cast(av_tiles), av_wg, "AV"); - ComputeOutParams p = make_compute_out_params(S, Hq, Hkv, D, context_len, g); - WGPUBuffer ubuf = graph.make_uniform_buffer(&p, sizeof(p)); - BufferBinding bindings[3] = { - {out.buffer, out.nbytes}, - {attn_weights_softmax, aw_bytes}, - {v_cache.buffer, v_cache.nbytes}}; - const char* av_src = kSdpaComputeOutWGSL; - if (graph.kv_f16()) { - av_src = kSdpaComputeOutHalfWGSL; - } - build_dispatch( + av_buf = graph.create_params_buffer(initial_state.compute_out); + const std::vector av_bindings = { + {out.buffer, 0u, out.nbytes}, + {attn_weights_softmax, 0u, aw_bytes}, + {v_cache.buffer, 0u, v_cache.nbytes}}; + add_sdpa_compute_dispatch( graph, - av_src, - bindings, - 3, - ubuf, - sizeof(p), - wgc.x, - wgc.y, + graph.kv_f16() ? kComputeOutHalfShader : kComputeOutShader, + av_bindings, + av_buf, + sizeof(ComputeOutParams), + initial_state.av_grid, av_wg, - true, "sdpa_compute_out"); - av_buf = ubuf; av_idx = graph.num_dispatches() - 1; + materialized_range.end = graph.num_dispatches(); } - // Per-step recompute: live S (q resize) or input_pos (SymInt); inert if - // static. - const int64_t pos_const = input_pos; - auto sdpa_resize = [q_id, - qn, - S, - out_id, - dynamic_pos, - input_pos_id, - pos_const, - Hq, - Hkv, - D, - Cmax, - g, - scale, - qk_idx, - uc_k_idx, - uc_v_idx, - softmax_idx, - av_idx, - uc_wg, - qk_wg, - av_wg, - uc_k_buf, - uc_v_buf, - qk_buf, - softmax_buf, - av_buf](WebGPUGraph& gr) { - const int64_t s = gr.cur_dims(q_id)[qn - 3]; - const int64_t pos = dynamic_pos - ? static_cast(gr.read_symint(input_pos_id)) - : pos_const; - if (s <= 0 || pos < 0) { - throw std::runtime_error("WebGPU sdpa: invalid live S or input_pos"); - } - // Scratch (attn_weights/softmax) is sized at build for S=max; a larger live - // S would overrun it. Make that invariant load-bearing. - if (s > S) { - throw std::runtime_error( - "WebGPU sdpa: live S exceeds the build-time max (scratch capacity)"); - } - const int64_t ctx = s + pos; - if (ctx <= 0 || ctx > Cmax) { - throw std::runtime_error( - "WebGPU sdpa: context_len exceeds cache capacity"); + WGPUBuffer k16_buf = nullptr; + size_t k16_idx = 0; + utils::DispatchRange k16_range = {}; + if (record_k16) { + k16_range.begin = graph.num_dispatches(); + k16_buf = graph.create_params_buffer(initial_state.streaming_k16); + const std::vector k16_bindings = { + {out.buffer, 0u, out.nbytes}, + {q.buffer, 0u, q.nbytes}, + {k_cache.buffer, 0u, k_cache.nbytes}, + {v_cache.buffer, 0u, v_cache.nbytes}}; + const utils::WgCount initial_grid = initial_state.use_fd + ? utils::WgCount{0u, 0u} + : initial_state.streaming_k16_grid; + add_sdpa_compute_dispatch( + graph, + k16_shader, + k16_bindings, + k16_buf, + sizeof(StreamingAttentionK16Params), + initial_grid, + 0, + k16_label); + k16_idx = graph.num_dispatches() - 1; + k16_range.end = graph.num_dispatches(); + } + + SdpaFdDecodeResources fd_resources = {}; + size_t route_group = 0; + if (record_fd) { + fd_resources = record_sdpa_fd_decode_dispatches( + graph, q, k_cache, v_cache, out, initial_state.fd); + } + if (dual_route) { + const utils::DispatchRange prefill_range = + record_k16 ? k16_range : materialized_range; + route_group = graph.register_dispatch_route_group( + {prefill_range, fd_resources.dispatch_range}); + } + + auto refresh_state = [compute_live_state, + q_id, + out_id, + dual_route, + record_k16, + record_materialized, + record_fd, + fixed_use_fd = initial_state.use_fd, + route_group, + uc_k_idx, + uc_v_idx, + qk_idx, + softmax_idx, + av_idx, + k16_idx, + uc_k_buf, + uc_v_buf, + qk_buf, + softmax_buf, + av_buf, + k16_buf, + fd_resources](WebGPUGraph& gr) { + const SdpaLiveState state = compute_live_state(gr); + + wgpuQueueWriteBuffer( + gr.queue(), + uc_k_buf, + 0, + &state.update_cache, + sizeof(state.update_cache)); + wgpuQueueWriteBuffer( + gr.queue(), + uc_v_buf, + 0, + &state.update_cache, + sizeof(state.update_cache)); + if (record_materialized) { + wgpuQueueWriteBuffer( + gr.queue(), + qk_buf, + 0, + &state.attn_weights, + sizeof(state.attn_weights)); + wgpuQueueWriteBuffer( + gr.queue(), softmax_buf, 0, &state.softmax, sizeof(state.softmax)); + wgpuQueueWriteBuffer( + gr.queue(), av_buf, 0, &state.compute_out, sizeof(state.compute_out)); } - const uint32_t kv_off = static_cast( - static_cast(pos) * static_cast(Hkv) * - static_cast(D)); - const uint64_t aw_floats = static_cast(Hq) * - static_cast(s) * static_cast(ctx); - if (aw_floats > UINT32_MAX) { - throw std::runtime_error("WebGPU sdpa: Hq*S*context_len exceeds uint32"); + if (record_k16) { + wgpuQueueWriteBuffer( + gr.queue(), + k16_buf, + 0, + &state.streaming_k16, + sizeof(state.streaming_k16)); } - const uint64_t kv_numel = static_cast(s) * - static_cast(Hkv) * static_cast(D); - if (kv_numel > UINT32_MAX) { - throw std::runtime_error("WebGPU sdpa: S*Hkv*D exceeds uint32"); + if (record_fd) { + write_sdpa_fd_decode_uniforms(gr.queue(), fd_resources, state.fd); } - const uint64_t k_cache_numel = static_cast(Cmax) * - static_cast(Hkv) * static_cast(D); - // update_cache K/V: dispatch (kv_numel) + dst offset scale with live S/pos. - UpdateCacheParams uc = - make_update_cache_params(kv_numel, kv_off, k_cache_numel); - wgpuQueueWriteBuffer(gr.queue(), uc_k_buf, 0, &uc, sizeof(uc)); - wgpuQueueWriteBuffer(gr.queue(), uc_v_buf, 0, &uc, sizeof(uc)); - const uint32_t uc_wgc = utils::compute_1d_workgroup_count( - gr.device(), static_cast(kv_numel), uc_wg, "uc(resize)"); - gr.dispatch_at(uc_k_idx).workgroup_count_x = uc_wgc; - gr.dispatch_at(uc_v_idx).workgroup_count_x = uc_wgc; - - // QK: one thread per TM x TN tile; grid = Hq*ceil(S/TM)*ceil(ctx/TN). - AttnWeightsParams qp = - make_attn_weights_params(s, Hq, Hkv, D, ctx, pos, g, scale); - wgpuQueueWriteBuffer(gr.queue(), qk_buf, 0, &qp, sizeof(qp)); - const int64_t qk_tiles = - Hq * utils::div_up(s, kSdpaTileM) * utils::div_up(ctx, kSdpaTileN); - const utils::WgCount qk_wgc = utils::compute_2d_workgroup_count( - gr.device(), static_cast(qk_tiles), qk_wg, "QK(resize)"); - gr.dispatch_at(qk_idx).workgroup_count_x = qk_wgc.x; - gr.dispatch_at(qk_idx).workgroup_count_y = qk_wgc.y; - - // softmax: one workgroup per (h,s) row. - SoftmaxParams sp = make_softmax_params(Hq, s, ctx); - wgpuQueueWriteBuffer(gr.queue(), softmax_buf, 0, &sp, sizeof(sp)); - const utils::WgCount sm_wgc = utils::compute_2d_workgroup_count( - gr.device(), static_cast(Hq * s), 1, "softmax(resize)"); - gr.dispatch_at(softmax_idx).workgroup_count_x = sm_wgc.x; - gr.dispatch_at(softmax_idx).workgroup_count_y = sm_wgc.y; - - // AV: one thread per TM x TN tile; grid = Hq*ceil(S/TM)*ceil(D/TN). - ComputeOutParams op = make_compute_out_params(s, Hq, Hkv, D, ctx, g); - wgpuQueueWriteBuffer(gr.queue(), av_buf, 0, &op, sizeof(op)); - const int64_t av_tiles = - Hq * utils::div_up(s, kSdpaTileM) * utils::div_up(D, kSdpaTileN); - const utils::WgCount av_wgc = utils::compute_2d_workgroup_count( - gr.device(), static_cast(av_tiles), av_wg, "AV(resize)"); - gr.dispatch_at(av_idx).workgroup_count_x = av_wgc.x; - gr.dispatch_at(av_idx).workgroup_count_y = av_wgc.y; - - // Output attn has the same shape as q: [.., S, Hq, D]. + gr.dispatch_at(uc_k_idx).workgroup_count_x = state.update_cache_grid.x; + gr.dispatch_at(uc_k_idx).workgroup_count_y = state.update_cache_grid.y; + gr.dispatch_at(uc_v_idx).workgroup_count_x = state.update_cache_grid.x; + gr.dispatch_at(uc_v_idx).workgroup_count_y = state.update_cache_grid.y; + if (dual_route) { + const size_t active_route = state.use_fd ? 1 : 0; + const std::vector active_grids = state.use_fd + ? std::vector< + utils::WgCount>{state.fd.split_grid, state.fd.reduce_grid} + : (record_k16 + ? std::vector{state.streaming_k16_grid} + : std::vector{ + state.qk_grid, state.softmax_grid, state.av_grid}); + gr.select_dispatch_route(route_group, active_route, active_grids); + } else if (state.use_fd) { + if (!fixed_use_fd) { + throw std::runtime_error("WebGPU sdpa: static route changed"); + } + gr.dispatch_at(fd_resources.dispatch_range.begin).workgroup_count_x = + state.fd.split_grid.x; + gr.dispatch_at(fd_resources.dispatch_range.begin).workgroup_count_y = + state.fd.split_grid.y; + gr.dispatch_at(fd_resources.dispatch_range.begin + 1).workgroup_count_x = + state.fd.reduce_grid.x; + gr.dispatch_at(fd_resources.dispatch_range.begin + 1).workgroup_count_y = + state.fd.reduce_grid.y; + } else if (record_k16) { + if (fixed_use_fd) { + throw std::runtime_error("WebGPU sdpa: static route changed"); + } + gr.dispatch_at(k16_idx).workgroup_count_x = state.streaming_k16_grid.x; + gr.dispatch_at(k16_idx).workgroup_count_y = state.streaming_k16_grid.y; + } else { + if (fixed_use_fd) { + throw std::runtime_error("WebGPU sdpa: static route changed"); + } + gr.dispatch_at(qk_idx).workgroup_count_x = state.qk_grid.x; + gr.dispatch_at(qk_idx).workgroup_count_y = state.qk_grid.y; + gr.dispatch_at(softmax_idx).workgroup_count_x = state.softmax_grid.x; + gr.dispatch_at(softmax_idx).workgroup_count_y = state.softmax_grid.y; + gr.dispatch_at(av_idx).workgroup_count_x = state.av_grid.x; + gr.dispatch_at(av_idx).workgroup_count_y = state.av_grid.y; + } gr.set_cur_dims(out_id, gr.cur_dims(q_id)); }; - // q and input_pos share one idempotent recompute; a double-fire is harmless. - graph.add_tensor_resize_hook(q_id, sdpa_resize); + + refresh_state(graph); + graph.add_tensor_resize_hook(q_id, refresh_state); if (dynamic_pos) { - graph.add_resize_hook(input_pos_id, sdpa_resize); + graph.add_resize_hook(input_pos_id, refresh_state); } } diff --git a/backends/webgpu/runtime/ops/sdpa/streaming_attention_k16_causal_bound.wgsl b/backends/webgpu/runtime/ops/sdpa/streaming_attention_k16_causal_bound.wgsl new file mode 100644 index 00000000000..5b8cf1de03c --- /dev/null +++ b/backends/webgpu/runtime/ops/sdpa/streaming_attention_k16_causal_bound.wgsl @@ -0,0 +1,268 @@ +enable f16; + +@group(0) @binding(0) var t_out: array>; +@group(0) @binding(1) var t_q: array>; +@group(0) @binding(2) var t_k_cache: array>; +@group(0) @binding(3) var t_v_cache: array>; + +struct Params { + S: u32, + context_len: u32, + input_pos: u32, + q_token_stride4: u32, + q_head_stride4: u32, + kv_token_stride4: u32, + kv_head_stride4: u32, + o_token_stride4: u32, + o_head_stride4: u32, + _pad0: u32, + _pad1: u32, + _pad2: u32, +} +@group(0) @binding(4) var params: Params; + +const HQ: u32 = 32u; +const HKV: u32 = 8u; +const G: u32 = 4u; +const D: u32 = 64u; +const D4: u32 = 16u; +const Q_TILE: u32 = 32u; +const K_TILE: u32 = 16u; +const SCALE: f32 = 0.125; +const NEG_INF: f32 = -1.0e30; + +var t_q_tile: array, 512>; +var t_k_tile: array, 256>; +var t_v_tile: array, 256>; +var t_scores: array, 128>; +var t_m: array; +var t_d: array; +var t_alpha: array; + +fn dot_qk(row: u32, key: u32) -> f32 { + let q_base = row * D4; + let k_base = key * D4; + var sum = 0.0; + var d4 = 0u; + loop { + if (d4 >= D4) { + break; + } + sum += dot(t_q_tile[q_base + d4], vec4(t_k_tile[k_base + d4])); + d4 += 1u; + } + return sum * SCALE; +} + +fn score_for( + row: u32, + key_in_tile: u32, + key: u32, + row_valid: bool, + key_valid: bool, + token: u32, +) -> f32 { + if (row_valid && key_valid && key <= params.input_pos + token) { + return dot_qk(row, key_in_tile); + } + return NEG_INF; +} + +fn max4(v: vec4) -> f32 { + return max(max(v.x, v.y), max(v.z, v.w)); +} + +fn exp_sum(v: vec4, maximum: f32) -> f32 { + let p = exp(v - vec4(maximum)); + return p.x + p.y + p.z + p.w; +} + +@compute @workgroup_size(32, 4, 1) +fn main( + @builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3, +) { + let groups_per_kv: u32 = (params.S * G + 31u) / 32u; + let kv_head = wid.x / groups_per_kv; + let row_group = wid.x % groups_per_kv; + if (kv_head >= HKV) { + return; + } + + let row: u32 = lid.x; + let dim_vec4_base: u32 = lid.y * 4u; + let logical_row: u32 = row_group * Q_TILE + row; + let row_valid = logical_row < params.S * G; + let token: u32 = logical_row / G; + let q_head: u32 = kv_head * G + logical_row % G; + let local_linear = lid.y * 32u + lid.x; + let group_max_logical_row = + min(params.S * G - 1u, row_group * Q_TILE + Q_TILE - 1u); + let group_max_token = group_max_logical_row / G; + let group_context_len = + min(params.context_len, params.input_pos + group_max_token + 1u); + + var load_slot = 0u; + loop { + if (load_slot >= 4u) { + break; + } + let tile_index = local_linear + load_slot * 128u; + let load_row = tile_index / D4; + let load_dim4 = tile_index % D4; + let load_logical_row = row_group * Q_TILE + load_row; + if (load_logical_row < params.S * G) { + let load_token = load_logical_row / G; + let load_q_head = kv_head * G + load_logical_row % G; + let q_index = + load_token * params.q_token_stride4 + + load_q_head * params.q_head_stride4 + + load_dim4; + t_q_tile[tile_index] = t_q[q_index]; + } else { + t_q_tile[tile_index] = vec4(0.0); + } + load_slot += 1u; + } + + if (lid.y == 0u) { + t_m[row] = NEG_INF; + t_d[row] = 0.0; + t_alpha[row] = 0.0; + } + workgroupBarrier(); + + var score_acc: vec4; + var output_acc: array, 4>; + score_acc = vec4(0.0); + output_acc[0] = vec4(0.0); + output_acc[1] = vec4(0.0); + output_acc[2] = vec4(0.0); + output_acc[3] = vec4(0.0); + + var key_tile_start = 0u; + loop { + if (key_tile_start >= group_context_len) { + break; + } + + load_slot = 0u; + loop { + if (load_slot >= 2u) { + break; + } + let tile_index = local_linear + load_slot * 128u; + let key_in_tile = tile_index / D4; + let load_dim4 = tile_index % D4; + let key = key_tile_start + key_in_tile; + if (key < params.context_len) { + let cache_index = + key * params.kv_token_stride4 + + kv_head * params.kv_head_stride4 + + load_dim4; + t_k_tile[tile_index] = t_k_cache[cache_index]; + t_v_tile[tile_index] = t_v_cache[cache_index]; + } else { + t_k_tile[tile_index] = vec4(0.0h); + t_v_tile[tile_index] = vec4(0.0h); + } + load_slot += 1u; + } + workgroupBarrier(); + + let score_key_base = lid.y * 4u; + let key0 = key_tile_start + score_key_base; + let key1 = key0 + 1u; + let key2 = key0 + 2u; + let key3 = key0 + 3u; + score_acc = vec4( + score_for(row, score_key_base, key0, row_valid, key0 < params.context_len, token), + score_for(row, score_key_base + 1u, key1, row_valid, key1 < params.context_len, token), + score_for(row, score_key_base + 2u, key2, row_valid, key2 < params.context_len, token), + score_for(row, score_key_base + 3u, key3, row_valid, key3 < params.context_len, token), + ); + let score_store = row * 4u + lid.y; + t_scores[score_store] = score_acc; + workgroupBarrier(); + + if (lid.y == 0u) { + let row_score_base = row * 4u; + let s0 = t_scores[row_score_base]; + let s1 = t_scores[row_score_base + 1u]; + let s2 = t_scores[row_score_base + 2u]; + let s3 = t_scores[row_score_base + 3u]; + let tile_max = max(max(max4(s0), max4(s1)), max(max4(s2), max4(s3))); + let old_m = t_m[row]; + let old_d = t_d[row]; + let new_m = max(old_m, tile_max); + if (row_valid) { + t_alpha[row] = exp(old_m - new_m); + let tile_sum = + exp_sum(s0, new_m) + exp_sum(s1, new_m) + + exp_sum(s2, new_m) + exp_sum(s3, new_m); + t_d[row] = old_d * t_alpha[row] + tile_sum; + t_m[row] = new_m; + } else { + t_alpha[row] = 0.0; + t_d[row] = 1.0; + t_m[row] = 0.0; + } + } + workgroupBarrier(); + + let alpha = t_alpha[row]; + output_acc[0] = output_acc[0] * alpha; + output_acc[1] = output_acc[1] * alpha; + output_acc[2] = output_acc[2] * alpha; + output_acc[3] = output_acc[3] * alpha; + let new_m = t_m[row]; + let row_score_base = row * 4u; + var score_block = 0u; + loop { + if (score_block >= 4u) { + break; + } + let probabilities = exp(t_scores[row_score_base + score_block] - vec4(new_m)); + let value_key_base = score_block * 4u; + let value_dim0 = dim_vec4_base; + let value_dim1 = dim_vec4_base + 1u; + let value_dim2 = dim_vec4_base + 2u; + let value_dim3 = dim_vec4_base + 3u; + output_acc[0] += + vec4(t_v_tile[(value_key_base + 0u) * D4 + value_dim0]) * probabilities.x + + vec4(t_v_tile[(value_key_base + 1u) * D4 + value_dim0]) * probabilities.y + + vec4(t_v_tile[(value_key_base + 2u) * D4 + value_dim0]) * probabilities.z + + vec4(t_v_tile[(value_key_base + 3u) * D4 + value_dim0]) * probabilities.w; + output_acc[1] += + vec4(t_v_tile[(value_key_base + 0u) * D4 + value_dim1]) * probabilities.x + + vec4(t_v_tile[(value_key_base + 1u) * D4 + value_dim1]) * probabilities.y + + vec4(t_v_tile[(value_key_base + 2u) * D4 + value_dim1]) * probabilities.z + + vec4(t_v_tile[(value_key_base + 3u) * D4 + value_dim1]) * probabilities.w; + output_acc[2] += + vec4(t_v_tile[(value_key_base + 0u) * D4 + value_dim2]) * probabilities.x + + vec4(t_v_tile[(value_key_base + 1u) * D4 + value_dim2]) * probabilities.y + + vec4(t_v_tile[(value_key_base + 2u) * D4 + value_dim2]) * probabilities.z + + vec4(t_v_tile[(value_key_base + 3u) * D4 + value_dim2]) * probabilities.w; + output_acc[3] += + vec4(t_v_tile[(value_key_base + 0u) * D4 + value_dim3]) * probabilities.x + + vec4(t_v_tile[(value_key_base + 1u) * D4 + value_dim3]) * probabilities.y + + vec4(t_v_tile[(value_key_base + 2u) * D4 + value_dim3]) * probabilities.z + + vec4(t_v_tile[(value_key_base + 3u) * D4 + value_dim3]) * probabilities.w; + score_block += 1u; + } + workgroupBarrier(); + key_tile_start += K_TILE; + } + + if (row_valid) { + let denominator = t_d[row]; + let output_base = + token * params.o_token_stride4 + + q_head * params.o_head_stride4 + + dim_vec4_base; + t_out[output_base] = output_acc[0] / denominator; + t_out[output_base + 1u] = output_acc[1] / denominator; + t_out[output_base + 2u] = output_acc[2] / denominator; + t_out[output_base + 3u] = output_acc[3] / denominator; + } +} diff --git a/backends/webgpu/runtime/ops/sdpa/streaming_attention_k16_causal_bound_wgsl.h b/backends/webgpu/runtime/ops/sdpa/streaming_attention_k16_causal_bound_wgsl.h new file mode 100644 index 00000000000..255c6afc6e0 --- /dev/null +++ b/backends/webgpu/runtime/ops/sdpa/streaming_attention_k16_causal_bound_wgsl.h @@ -0,0 +1,292 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from streaming_attention_k16_causal_bound.wgsl - DO NOT EDIT. +// wgsl-sha256: b1435c4f72834cb896eb4248a899aedeabbdcc772818afc419feca03e8957ffd +inline constexpr const char* kStreamingAttentionK16CausalBoundWGSL = R"( +enable f16; + +@group(0) @binding(0) var t_out: array>; +@group(0) @binding(1) var t_q: array>; +@group(0) @binding(2) var t_k_cache: array>; +@group(0) @binding(3) var t_v_cache: array>; + +struct Params { + S: u32, + context_len: u32, + input_pos: u32, + q_token_stride4: u32, + q_head_stride4: u32, + kv_token_stride4: u32, + kv_head_stride4: u32, + o_token_stride4: u32, + o_head_stride4: u32, + _pad0: u32, + _pad1: u32, + _pad2: u32, +} +@group(0) @binding(4) var params: Params; + +const HQ: u32 = 32u; +const HKV: u32 = 8u; +const G: u32 = 4u; +const D: u32 = 64u; +const D4: u32 = 16u; +const Q_TILE: u32 = 32u; +const K_TILE: u32 = 16u; +const SCALE: f32 = 0.125; +const NEG_INF: f32 = -1.0e30; + +var t_q_tile: array, 512>; +var t_k_tile: array, 256>; +var t_v_tile: array, 256>; +var t_scores: array, 128>; +var t_m: array; +var t_d: array; +var t_alpha: array; + +fn dot_qk(row: u32, key: u32) -> f32 { + let q_base = row * D4; + let k_base = key * D4; + var sum = 0.0; + var d4 = 0u; + loop { + if (d4 >= D4) { + break; + } + sum += dot(t_q_tile[q_base + d4], vec4(t_k_tile[k_base + d4])); + d4 += 1u; + } + return sum * SCALE; +} + +fn score_for( + row: u32, + key_in_tile: u32, + key: u32, + row_valid: bool, + key_valid: bool, + token: u32, +) -> f32 { + if (row_valid && key_valid && key <= params.input_pos + token) { + return dot_qk(row, key_in_tile); + } + return NEG_INF; +} + +fn max4(v: vec4) -> f32 { + return max(max(v.x, v.y), max(v.z, v.w)); +} + +fn exp_sum(v: vec4, maximum: f32) -> f32 { + let p = exp(v - vec4(maximum)); + return p.x + p.y + p.z + p.w; +} + +@compute @workgroup_size(32, 4, 1) +fn main( + @builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3, +) { + let groups_per_kv: u32 = (params.S * G + 31u) / 32u; + let kv_head = wid.x / groups_per_kv; + let row_group = wid.x % groups_per_kv; + if (kv_head >= HKV) { + return; + } + + let row: u32 = lid.x; + let dim_vec4_base: u32 = lid.y * 4u; + let logical_row: u32 = row_group * Q_TILE + row; + let row_valid = logical_row < params.S * G; + let token: u32 = logical_row / G; + let q_head: u32 = kv_head * G + logical_row % G; + let local_linear = lid.y * 32u + lid.x; + let group_max_logical_row = + min(params.S * G - 1u, row_group * Q_TILE + Q_TILE - 1u); + let group_max_token = group_max_logical_row / G; + let group_context_len = + min(params.context_len, params.input_pos + group_max_token + 1u); + + var load_slot = 0u; + loop { + if (load_slot >= 4u) { + break; + } + let tile_index = local_linear + load_slot * 128u; + let load_row = tile_index / D4; + let load_dim4 = tile_index % D4; + let load_logical_row = row_group * Q_TILE + load_row; + if (load_logical_row < params.S * G) { + let load_token = load_logical_row / G; + let load_q_head = kv_head * G + load_logical_row % G; + let q_index = + load_token * params.q_token_stride4 + + load_q_head * params.q_head_stride4 + + load_dim4; + t_q_tile[tile_index] = t_q[q_index]; + } else { + t_q_tile[tile_index] = vec4(0.0); + } + load_slot += 1u; + } + + if (lid.y == 0u) { + t_m[row] = NEG_INF; + t_d[row] = 0.0; + t_alpha[row] = 0.0; + } + workgroupBarrier(); + + var score_acc: vec4; + var output_acc: array, 4>; + score_acc = vec4(0.0); + output_acc[0] = vec4(0.0); + output_acc[1] = vec4(0.0); + output_acc[2] = vec4(0.0); + output_acc[3] = vec4(0.0); + + var key_tile_start = 0u; + loop { + if (key_tile_start >= group_context_len) { + break; + } + + load_slot = 0u; + loop { + if (load_slot >= 2u) { + break; + } + let tile_index = local_linear + load_slot * 128u; + let key_in_tile = tile_index / D4; + let load_dim4 = tile_index % D4; + let key = key_tile_start + key_in_tile; + if (key < params.context_len) { + let cache_index = + key * params.kv_token_stride4 + + kv_head * params.kv_head_stride4 + + load_dim4; + t_k_tile[tile_index] = t_k_cache[cache_index]; + t_v_tile[tile_index] = t_v_cache[cache_index]; + } else { + t_k_tile[tile_index] = vec4(0.0h); + t_v_tile[tile_index] = vec4(0.0h); + } + load_slot += 1u; + } + workgroupBarrier(); + + let score_key_base = lid.y * 4u; + let key0 = key_tile_start + score_key_base; + let key1 = key0 + 1u; + let key2 = key0 + 2u; + let key3 = key0 + 3u; + score_acc = vec4( + score_for(row, score_key_base, key0, row_valid, key0 < params.context_len, token), + score_for(row, score_key_base + 1u, key1, row_valid, key1 < params.context_len, token), + score_for(row, score_key_base + 2u, key2, row_valid, key2 < params.context_len, token), + score_for(row, score_key_base + 3u, key3, row_valid, key3 < params.context_len, token), + ); + let score_store = row * 4u + lid.y; + t_scores[score_store] = score_acc; + workgroupBarrier(); + + if (lid.y == 0u) { + let row_score_base = row * 4u; + let s0 = t_scores[row_score_base]; + let s1 = t_scores[row_score_base + 1u]; + let s2 = t_scores[row_score_base + 2u]; + let s3 = t_scores[row_score_base + 3u]; + let tile_max = max(max(max4(s0), max4(s1)), max(max4(s2), max4(s3))); + let old_m = t_m[row]; + let old_d = t_d[row]; + let new_m = max(old_m, tile_max); + if (row_valid) { + t_alpha[row] = exp(old_m - new_m); + let tile_sum = + exp_sum(s0, new_m) + exp_sum(s1, new_m) + + exp_sum(s2, new_m) + exp_sum(s3, new_m); + t_d[row] = old_d * t_alpha[row] + tile_sum; + t_m[row] = new_m; + } else { + t_alpha[row] = 0.0; + t_d[row] = 1.0; + t_m[row] = 0.0; + } + } + workgroupBarrier(); + + let alpha = t_alpha[row]; + output_acc[0] = output_acc[0] * alpha; + output_acc[1] = output_acc[1] * alpha; + output_acc[2] = output_acc[2] * alpha; + output_acc[3] = output_acc[3] * alpha; + let new_m = t_m[row]; + let row_score_base = row * 4u; + var score_block = 0u; + loop { + if (score_block >= 4u) { + break; + } + let probabilities = exp(t_scores[row_score_base + score_block] - vec4(new_m)); + let value_key_base = score_block * 4u; + let value_dim0 = dim_vec4_base; + let value_dim1 = dim_vec4_base + 1u; + let value_dim2 = dim_vec4_base + 2u; + let value_dim3 = dim_vec4_base + 3u; + output_acc[0] += + vec4(t_v_tile[(value_key_base + 0u) * D4 + value_dim0]) * probabilities.x + + vec4(t_v_tile[(value_key_base + 1u) * D4 + value_dim0]) * probabilities.y + + vec4(t_v_tile[(value_key_base + 2u) * D4 + value_dim0]) * probabilities.z + + vec4(t_v_tile[(value_key_base + 3u) * D4 + value_dim0]) * probabilities.w; + output_acc[1] += + vec4(t_v_tile[(value_key_base + 0u) * D4 + value_dim1]) * probabilities.x + + vec4(t_v_tile[(value_key_base + 1u) * D4 + value_dim1]) * probabilities.y + + vec4(t_v_tile[(value_key_base + 2u) * D4 + value_dim1]) * probabilities.z + + vec4(t_v_tile[(value_key_base + 3u) * D4 + value_dim1]) * probabilities.w; + output_acc[2] += + vec4(t_v_tile[(value_key_base + 0u) * D4 + value_dim2]) * probabilities.x + + vec4(t_v_tile[(value_key_base + 1u) * D4 + value_dim2]) * probabilities.y + + vec4(t_v_tile[(value_key_base + 2u) * D4 + value_dim2]) * probabilities.z + + vec4(t_v_tile[(value_key_base + 3u) * D4 + value_dim2]) * probabilities.w; + output_acc[3] += + vec4(t_v_tile[(value_key_base + 0u) * D4 + value_dim3]) * probabilities.x + + vec4(t_v_tile[(value_key_base + 1u) * D4 + value_dim3]) * probabilities.y + + vec4(t_v_tile[(value_key_base + 2u) * D4 + value_dim3]) * probabilities.z + + vec4(t_v_tile[(value_key_base + 3u) * D4 + value_dim3]) * probabilities.w; + score_block += 1u; + } + workgroupBarrier(); + key_tile_start += K_TILE; + } + + if (row_valid) { + let denominator = t_d[row]; + let output_base = + token * params.o_token_stride4 + + q_head * params.o_head_stride4 + + dim_vec4_base; + t_out[output_base] = output_acc[0] / denominator; + t_out[output_base + 1u] = output_acc[1] / denominator; + t_out[output_base + 2u] = output_acc[2] / denominator; + t_out[output_base + 3u] = output_acc[3] / denominator; + } +} +)"; + +inline constexpr uint32_t kStreamingAttentionK16CausalBoundWorkgroupSizeX = 32; +inline constexpr uint32_t kStreamingAttentionK16CausalBoundWorkgroupSizeY = 4; +inline constexpr uint32_t kStreamingAttentionK16CausalBoundWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_k16_causal_bound.wgsl b/backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_k16_causal_bound.wgsl new file mode 100644 index 00000000000..5e27346af6f --- /dev/null +++ b/backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_k16_causal_bound.wgsl @@ -0,0 +1,283 @@ +enable f16; + +@group(0) @binding(0) var t_out: array>; +@group(0) @binding(1) var t_q: array>; +@group(0) @binding(2) var t_k_cache: array>; +@group(0) @binding(3) var t_v_cache: array>; + +struct Params { + S: u32, + context_len: u32, + input_pos: u32, + q_token_stride4: u32, + q_head_stride4: u32, + kv_token_stride4: u32, + kv_head_stride4: u32, + o_token_stride4: u32, + o_head_stride4: u32, + _pad0: u32, + _pad1: u32, + _pad2: u32, +} +@group(0) @binding(4) var params: Params; + +const HQ: u32 = 16u; +const HKV: u32 = 8u; +const G: u32 = 2u; +const D: u32 = 128u; +const D4: u32 = 32u; +const Q_TILE: u32 = 16u; +const K_TILE: u32 = 16u; +const SCALE: f32 = 0.08838834764831845; +const NEG_INF: f32 = -1.0e30; + +var t_q_tile: array, 512>; +var t_kv_tile: array, 512>; +var t_scores: array, 128>; +var t_m: array; +var t_d: array; +var t_alpha: array; + +fn dot_qk(row: u32, key: u32) -> f32 { + let q_base = row * D4; + let k_base = key * D4; + var sum = 0.0; + var d4 = 0u; + loop { + if (d4 >= D4) { + break; + } + sum += dot(t_q_tile[q_base + d4], vec4(t_kv_tile[k_base + d4])); + d4 += 1u; + } + return sum * SCALE; +} + +fn score_for( + row: u32, + key_in_tile: u32, + key: u32, + row_valid: bool, + key_valid: bool, + token: u32, +) -> f32 { + if (row_valid && key_valid && key <= params.input_pos + token) { + return dot_qk(row, key_in_tile); + } + return NEG_INF; +} + +fn max2(v: vec2) -> f32 { + return max(v.x, v.y); +} + +fn exp_sum(v: vec2, maximum: f32) -> f32 { + let p = exp(v - vec2(maximum)); + return p.x + p.y; +} + +@compute @workgroup_size(16, 8, 1) +fn main( + @builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3, +) { + let groups_per_kv: u32 = (params.S * G + 15u) / 16u; + let kv_head = wid.x / groups_per_kv; + let row_group = wid.x % groups_per_kv; + if (kv_head >= HKV) { + return; + } + + let row: u32 = lid.x; + let dim_vec4_base: u32 = lid.y * 4u; + let logical_row: u32 = row_group * Q_TILE + row; + let row_valid = logical_row < params.S * G; + let token: u32 = logical_row / G; + let q_head: u32 = kv_head * G + logical_row % G; + let local_linear = lid.y * 16u + lid.x; + let group_max_logical_row = + min(params.S * G - 1u, row_group * Q_TILE + Q_TILE - 1u); + let group_max_token = group_max_logical_row / G; + let group_context_len = + min(params.context_len, params.input_pos + group_max_token + 1u); + + var load_slot = 0u; + loop { + if (load_slot >= 4u) { + break; + } + let tile_index = local_linear + load_slot * 128u; + let load_row = tile_index / D4; + let load_dim4 = tile_index % D4; + let load_logical_row = row_group * Q_TILE + load_row; + if (load_logical_row < params.S * G) { + let load_token = load_logical_row / G; + let load_q_head = kv_head * G + load_logical_row % G; + let q_index = + load_token * params.q_token_stride4 + + load_q_head * params.q_head_stride4 + + load_dim4; + t_q_tile[tile_index] = t_q[q_index]; + } else { + t_q_tile[tile_index] = vec4(0.0); + } + load_slot += 1u; + } + + if (lid.y == 0u) { + t_m[row] = NEG_INF; + t_d[row] = 0.0; + t_alpha[row] = 0.0; + } + workgroupBarrier(); + + var score_acc = vec2(0.0); + var output_acc: array, 4>; + output_acc[0] = vec4(0.0); + output_acc[1] = vec4(0.0); + output_acc[2] = vec4(0.0); + output_acc[3] = vec4(0.0); + + var key_tile_start = 0u; + loop { + if (key_tile_start >= group_context_len) { + break; + } + + load_slot = 0u; + loop { + if (load_slot >= 4u) { + break; + } + let tile_index = local_linear + load_slot * 128u; + let key_in_tile = tile_index / D4; + let load_dim4 = tile_index % D4; + let key = key_tile_start + key_in_tile; + if (key < params.context_len) { + let cache_index = + key * params.kv_token_stride4 + + kv_head * params.kv_head_stride4 + + load_dim4; + t_kv_tile[tile_index] = t_k_cache[cache_index]; + } else { + t_kv_tile[tile_index] = vec4(0.0h); + } + load_slot += 1u; + } + workgroupBarrier(); + + let score_key_base = lid.y * 2u; + let key0 = key_tile_start + score_key_base; + let key1 = key0 + 1u; + score_acc = vec2( + score_for(row, score_key_base, key0, row_valid, key0 < params.context_len, token), + score_for(row, score_key_base + 1u, key1, row_valid, key1 < params.context_len, token), + ); + let score_store = row * 8u + lid.y; + t_scores[score_store] = score_acc; + workgroupBarrier(); + + if (lid.y == 0u) { + let row_score_base = row * 8u; + let s0 = t_scores[row_score_base]; + let s1 = t_scores[row_score_base + 1u]; + let s2 = t_scores[row_score_base + 2u]; + let s3 = t_scores[row_score_base + 3u]; + let s4 = t_scores[row_score_base + 4u]; + let s5 = t_scores[row_score_base + 5u]; + let s6 = t_scores[row_score_base + 6u]; + let s7 = t_scores[row_score_base + 7u]; + let tile_max = max( + max(max(max2(s0), max2(s1)), max(max2(s2), max2(s3))), + max(max(max2(s4), max2(s5)), max(max2(s6), max2(s7))), + ); + let old_m = t_m[row]; + let old_d = t_d[row]; + let new_m = max(old_m, tile_max); + if (row_valid) { + t_alpha[row] = exp(old_m - new_m); + let tile_sum = + exp_sum(s0, new_m) + exp_sum(s1, new_m) + + exp_sum(s2, new_m) + exp_sum(s3, new_m) + + exp_sum(s4, new_m) + exp_sum(s5, new_m) + + exp_sum(s6, new_m) + exp_sum(s7, new_m); + t_d[row] = old_d * t_alpha[row] + tile_sum; + t_m[row] = new_m; + } else { + t_alpha[row] = 0.0; + t_d[row] = 1.0; + t_m[row] = 0.0; + } + } + + load_slot = 0u; + loop { + if (load_slot >= 4u) { + break; + } + let tile_index = local_linear + load_slot * 128u; + let key_in_tile = tile_index / D4; + let load_dim4 = tile_index % D4; + let key = key_tile_start + key_in_tile; + if (key < params.context_len) { + let cache_index = + key * params.kv_token_stride4 + + kv_head * params.kv_head_stride4 + + load_dim4; + t_kv_tile[tile_index] = t_v_cache[cache_index]; + } else { + t_kv_tile[tile_index] = vec4(0.0h); + } + load_slot += 1u; + } + workgroupBarrier(); + + let alpha = t_alpha[row]; + output_acc[0] = output_acc[0] * alpha; + output_acc[1] = output_acc[1] * alpha; + output_acc[2] = output_acc[2] * alpha; + output_acc[3] = output_acc[3] * alpha; + let new_m = t_m[row]; + let row_score_base = row * 8u; + var score_block = 0u; + loop { + if (score_block >= 8u) { + break; + } + let probabilities = + exp(t_scores[row_score_base + score_block] - vec2(new_m)); + let value_key_base = score_block * 2u; + let value_dim0 = dim_vec4_base; + let value_dim1 = dim_vec4_base + 1u; + let value_dim2 = dim_vec4_base + 2u; + let value_dim3 = dim_vec4_base + 3u; + output_acc[0] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim0]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim0]) * probabilities.y; + output_acc[1] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim1]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim1]) * probabilities.y; + output_acc[2] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim2]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim2]) * probabilities.y; + output_acc[3] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim3]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim3]) * probabilities.y; + score_block += 1u; + } + workgroupBarrier(); + key_tile_start += K_TILE; + } + + if (row_valid) { + let denominator = t_d[row]; + let output_base = + token * params.o_token_stride4 + + q_head * params.o_head_stride4 + + dim_vec4_base; + t_out[output_base] = output_acc[0] / denominator; + t_out[output_base + 1u] = output_acc[1] / denominator; + t_out[output_base + 2u] = output_acc[2] / denominator; + t_out[output_base + 3u] = output_acc[3] / denominator; + } +} diff --git a/backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_k16_causal_bound_wgsl.h b/backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_k16_causal_bound_wgsl.h new file mode 100644 index 00000000000..f69a4d72dae --- /dev/null +++ b/backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_k16_causal_bound_wgsl.h @@ -0,0 +1,311 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from streaming_attention_qwen3_k16_causal_bound.wgsl +// DO NOT EDIT. +// wgsl-sha256: b26e81d92ba58c833f7dfbf2da6bce1b9cb468d70d893cef9df3a7fae16ee1fc +inline constexpr const char* kStreamingAttentionQwen3K16CausalBoundWGSL = R"( +enable f16; + +@group(0) @binding(0) var t_out: array>; +@group(0) @binding(1) var t_q: array>; +@group(0) @binding(2) var t_k_cache: array>; +@group(0) @binding(3) var t_v_cache: array>; + +struct Params { + S: u32, + context_len: u32, + input_pos: u32, + q_token_stride4: u32, + q_head_stride4: u32, + kv_token_stride4: u32, + kv_head_stride4: u32, + o_token_stride4: u32, + o_head_stride4: u32, + _pad0: u32, + _pad1: u32, + _pad2: u32, +} +@group(0) @binding(4) var params: Params; + +const HQ: u32 = 16u; +const HKV: u32 = 8u; +const G: u32 = 2u; +const D: u32 = 128u; +const D4: u32 = 32u; +const Q_TILE: u32 = 16u; +const K_TILE: u32 = 16u; +const SCALE: f32 = 0.08838834764831845; +const NEG_INF: f32 = -1.0e30; + +var t_q_tile: array, 512>; +var t_kv_tile: array, 512>; +var t_scores: array, 128>; +var t_m: array; +var t_d: array; +var t_alpha: array; + +fn dot_qk(row: u32, key: u32) -> f32 { + let q_base = row * D4; + let k_base = key * D4; + var sum = 0.0; + var d4 = 0u; + loop { + if (d4 >= D4) { + break; + } + sum += dot(t_q_tile[q_base + d4], vec4(t_kv_tile[k_base + d4])); + d4 += 1u; + } + return sum * SCALE; +} + +fn score_for( + row: u32, + key_in_tile: u32, + key: u32, + row_valid: bool, + key_valid: bool, + token: u32, +) -> f32 { + if (row_valid && key_valid && key <= params.input_pos + token) { + return dot_qk(row, key_in_tile); + } + return NEG_INF; +} + +fn max2(v: vec2) -> f32 { + return max(v.x, v.y); +} + +fn exp_sum(v: vec2, maximum: f32) -> f32 { + let p = exp(v - vec2(maximum)); + return p.x + p.y; +} + +@compute @workgroup_size(16, 8, 1) +fn main( + @builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3, +) { + let groups_per_kv: u32 = (params.S * G + 15u) / 16u; + let kv_head = wid.x / groups_per_kv; + let row_group = wid.x % groups_per_kv; + if (kv_head >= HKV) { + return; + } + + let row: u32 = lid.x; + let dim_vec4_base: u32 = lid.y * 4u; + let logical_row: u32 = row_group * Q_TILE + row; + let row_valid = logical_row < params.S * G; + let token: u32 = logical_row / G; + let q_head: u32 = kv_head * G + logical_row % G; + let local_linear = lid.y * 16u + lid.x; + let group_max_logical_row = + min(params.S * G - 1u, row_group * Q_TILE + Q_TILE - 1u); + let group_max_token = group_max_logical_row / G; + let group_context_len = + min(params.context_len, params.input_pos + group_max_token + 1u); + + var load_slot = 0u; + loop { + if (load_slot >= 4u) { + break; + } + let tile_index = local_linear + load_slot * 128u; + let load_row = tile_index / D4; + let load_dim4 = tile_index % D4; + let load_logical_row = row_group * Q_TILE + load_row; + if (load_logical_row < params.S * G) { + let load_token = load_logical_row / G; + let load_q_head = kv_head * G + load_logical_row % G; + let q_index = + load_token * params.q_token_stride4 + + load_q_head * params.q_head_stride4 + + load_dim4; + t_q_tile[tile_index] = t_q[q_index]; + } else { + t_q_tile[tile_index] = vec4(0.0); + } + load_slot += 1u; + } + + if (lid.y == 0u) { + t_m[row] = NEG_INF; + t_d[row] = 0.0; + t_alpha[row] = 0.0; + } + workgroupBarrier(); + + var score_acc = vec2(0.0); + var output_acc: array, 4>; + output_acc[0] = vec4(0.0); + output_acc[1] = vec4(0.0); + output_acc[2] = vec4(0.0); + output_acc[3] = vec4(0.0); + + var key_tile_start = 0u; + loop { + if (key_tile_start >= group_context_len) { + break; + } + + load_slot = 0u; + loop { + if (load_slot >= 4u) { + break; + } + let tile_index = local_linear + load_slot * 128u; + let key_in_tile = tile_index / D4; + let load_dim4 = tile_index % D4; + let key = key_tile_start + key_in_tile; + if (key < params.context_len) { + let cache_index = + key * params.kv_token_stride4 + + kv_head * params.kv_head_stride4 + + load_dim4; + t_kv_tile[tile_index] = t_k_cache[cache_index]; + } else { + t_kv_tile[tile_index] = vec4(0.0h); + } + load_slot += 1u; + } + workgroupBarrier(); + + let score_key_base = lid.y * 2u; + let key0 = key_tile_start + score_key_base; + let key1 = key0 + 1u; + score_acc = vec2( + score_for(row, score_key_base, key0, row_valid, key0 < params.context_len, token), + score_for(row, score_key_base + 1u, key1, row_valid, key1 < params.context_len, token), + ); + let score_store = row * 8u + lid.y; + t_scores[score_store] = score_acc; + workgroupBarrier(); + + if (lid.y == 0u) { + let row_score_base = row * 8u; + let s0 = t_scores[row_score_base]; + let s1 = t_scores[row_score_base + 1u]; + let s2 = t_scores[row_score_base + 2u]; + let s3 = t_scores[row_score_base + 3u]; + let s4 = t_scores[row_score_base + 4u]; + let s5 = t_scores[row_score_base + 5u]; + let s6 = t_scores[row_score_base + 6u]; + let s7 = t_scores[row_score_base + 7u]; + let tile_max = max( + max(max(max2(s0), max2(s1)), max(max2(s2), max2(s3))), + max(max(max2(s4), max2(s5)), max(max2(s6), max2(s7))), + ); + let old_m = t_m[row]; + let old_d = t_d[row]; + let new_m = max(old_m, tile_max); + if (row_valid) { + t_alpha[row] = exp(old_m - new_m); + let tile_sum = + exp_sum(s0, new_m) + exp_sum(s1, new_m) + + exp_sum(s2, new_m) + exp_sum(s3, new_m) + + exp_sum(s4, new_m) + exp_sum(s5, new_m) + + exp_sum(s6, new_m) + exp_sum(s7, new_m); + t_d[row] = old_d * t_alpha[row] + tile_sum; + t_m[row] = new_m; + } else { + t_alpha[row] = 0.0; + t_d[row] = 1.0; + t_m[row] = 0.0; + } + } + + load_slot = 0u; + loop { + if (load_slot >= 4u) { + break; + } + let tile_index = local_linear + load_slot * 128u; + let key_in_tile = tile_index / D4; + let load_dim4 = tile_index % D4; + let key = key_tile_start + key_in_tile; + if (key < params.context_len) { + let cache_index = + key * params.kv_token_stride4 + + kv_head * params.kv_head_stride4 + + load_dim4; + t_kv_tile[tile_index] = t_v_cache[cache_index]; + } else { + t_kv_tile[tile_index] = vec4(0.0h); + } + load_slot += 1u; + } + workgroupBarrier(); + + let alpha = t_alpha[row]; + output_acc[0] = output_acc[0] * alpha; + output_acc[1] = output_acc[1] * alpha; + output_acc[2] = output_acc[2] * alpha; + output_acc[3] = output_acc[3] * alpha; + let new_m = t_m[row]; + let row_score_base = row * 8u; + var score_block = 0u; + loop { + if (score_block >= 8u) { + break; + } + let probabilities = + exp(t_scores[row_score_base + score_block] - vec2(new_m)); + let value_key_base = score_block * 2u; + let value_dim0 = dim_vec4_base; + let value_dim1 = dim_vec4_base + 1u; + let value_dim2 = dim_vec4_base + 2u; + let value_dim3 = dim_vec4_base + 3u; + output_acc[0] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim0]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim0]) * probabilities.y; + output_acc[1] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim1]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim1]) * probabilities.y; + output_acc[2] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim2]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim2]) * probabilities.y; + output_acc[3] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim3]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim3]) * probabilities.y; + score_block += 1u; + } + workgroupBarrier(); + key_tile_start += K_TILE; + } + + if (row_valid) { + let denominator = t_d[row]; + let output_base = + token * params.o_token_stride4 + + q_head * params.o_head_stride4 + + dim_vec4_base; + t_out[output_base] = output_acc[0] / denominator; + t_out[output_base + 1u] = output_acc[1] / denominator; + t_out[output_base + 2u] = output_acc[2] / denominator; + t_out[output_base + 3u] = output_acc[3] / denominator; + } +} +)"; + +inline constexpr uint32_t kStreamingAttentionQwen3K16CausalBoundWorkgroupSizeX = + 16; +inline constexpr uint32_t kStreamingAttentionQwen3K16CausalBoundWorkgroupSizeY = + 8; +inline constexpr uint32_t kStreamingAttentionQwen3K16CausalBoundWorkgroupSizeZ = + 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_q32_k16_causal_bound.wgsl b/backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_q32_k16_causal_bound.wgsl new file mode 100644 index 00000000000..6135b15b75d --- /dev/null +++ b/backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_q32_k16_causal_bound.wgsl @@ -0,0 +1,283 @@ +enable f16; + +@group(0) @binding(0) var t_out: array>; +@group(0) @binding(1) var t_q: array>; +@group(0) @binding(2) var t_k_cache: array>; +@group(0) @binding(3) var t_v_cache: array>; + +struct Params { + S: u32, + context_len: u32, + input_pos: u32, + q_token_stride4: u32, + q_head_stride4: u32, + kv_token_stride4: u32, + kv_head_stride4: u32, + o_token_stride4: u32, + o_head_stride4: u32, + _pad0: u32, + _pad1: u32, + _pad2: u32, +} +@group(0) @binding(4) var params: Params; + +const HQ: u32 = 16u; +const HKV: u32 = 8u; +const G: u32 = 2u; +const D: u32 = 128u; +const D4: u32 = 32u; +const Q_TILE: u32 = 32u; +const K_TILE: u32 = 16u; +const SCALE: f32 = 0.08838834764831845; +const NEG_INF: f32 = -1.0e30; + +var t_q_tile: array, 1024>; +var t_kv_tile: array, 512>; +var t_scores: array, 256>; +var t_m: array; +var t_d: array; +var t_alpha: array; + +fn dot_qk(row: u32, key: u32) -> f32 { + let q_base = row * D4; + let k_base = key * D4; + var sum = 0.0; + var d4 = 0u; + loop { + if (d4 >= D4) { + break; + } + sum += dot(t_q_tile[q_base + d4], vec4(t_kv_tile[k_base + d4])); + d4 += 1u; + } + return sum * SCALE; +} + +fn score_for( + row: u32, + key_in_tile: u32, + key: u32, + row_valid: bool, + key_valid: bool, + token: u32, +) -> f32 { + if (row_valid && key_valid && key <= params.input_pos + token) { + return dot_qk(row, key_in_tile); + } + return NEG_INF; +} + +fn max2(v: vec2) -> f32 { + return max(v.x, v.y); +} + +fn exp_sum(v: vec2, maximum: f32) -> f32 { + let p = exp(v - vec2(maximum)); + return p.x + p.y; +} + +@compute @workgroup_size(32, 8, 1) +fn main( + @builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3, +) { + let groups_per_kv: u32 = (params.S * G + 31u) / 32u; + let kv_head = wid.x / groups_per_kv; + let row_group = wid.x % groups_per_kv; + if (kv_head >= HKV) { + return; + } + + let row: u32 = lid.x; + let dim_vec4_base: u32 = lid.y * 4u; + let logical_row: u32 = row_group * Q_TILE + row; + let row_valid = logical_row < params.S * G; + let token: u32 = logical_row / G; + let q_head: u32 = kv_head * G + logical_row % G; + let local_linear = lid.y * 32u + lid.x; + let group_max_logical_row = + min(params.S * G - 1u, row_group * Q_TILE + Q_TILE - 1u); + let group_max_token = group_max_logical_row / G; + let group_context_len = + min(params.context_len, params.input_pos + group_max_token + 1u); + + var load_slot = 0u; + loop { + if (load_slot >= 4u) { + break; + } + let tile_index = local_linear + load_slot * 256u; + let load_row = tile_index / D4; + let load_dim4 = tile_index % D4; + let load_logical_row = row_group * Q_TILE + load_row; + if (load_logical_row < params.S * G) { + let load_token = load_logical_row / G; + let load_q_head = kv_head * G + load_logical_row % G; + let q_index = + load_token * params.q_token_stride4 + + load_q_head * params.q_head_stride4 + + load_dim4; + t_q_tile[tile_index] = t_q[q_index]; + } else { + t_q_tile[tile_index] = vec4(0.0); + } + load_slot += 1u; + } + + if (lid.y == 0u) { + t_m[row] = NEG_INF; + t_d[row] = 0.0; + t_alpha[row] = 0.0; + } + workgroupBarrier(); + + var score_acc = vec2(0.0); + var output_acc: array, 4>; + output_acc[0] = vec4(0.0); + output_acc[1] = vec4(0.0); + output_acc[2] = vec4(0.0); + output_acc[3] = vec4(0.0); + + var key_tile_start = 0u; + loop { + if (key_tile_start >= group_context_len) { + break; + } + + load_slot = 0u; + loop { + if (load_slot >= 2u) { + break; + } + let tile_index = local_linear + load_slot * 256u; + let key_in_tile = tile_index / D4; + let load_dim4 = tile_index % D4; + let key = key_tile_start + key_in_tile; + if (key < params.context_len) { + let cache_index = + key * params.kv_token_stride4 + + kv_head * params.kv_head_stride4 + + load_dim4; + t_kv_tile[tile_index] = t_k_cache[cache_index]; + } else { + t_kv_tile[tile_index] = vec4(0.0h); + } + load_slot += 1u; + } + workgroupBarrier(); + + let score_key_base = lid.y * 2u; + let key0 = key_tile_start + score_key_base; + let key1 = key0 + 1u; + score_acc = vec2( + score_for(row, score_key_base, key0, row_valid, key0 < params.context_len, token), + score_for(row, score_key_base + 1u, key1, row_valid, key1 < params.context_len, token), + ); + let score_store = row * 8u + lid.y; + t_scores[score_store] = score_acc; + workgroupBarrier(); + + if (lid.y == 0u) { + let row_score_base = row * 8u; + let s0 = t_scores[row_score_base]; + let s1 = t_scores[row_score_base + 1u]; + let s2 = t_scores[row_score_base + 2u]; + let s3 = t_scores[row_score_base + 3u]; + let s4 = t_scores[row_score_base + 4u]; + let s5 = t_scores[row_score_base + 5u]; + let s6 = t_scores[row_score_base + 6u]; + let s7 = t_scores[row_score_base + 7u]; + let tile_max = max( + max(max(max2(s0), max2(s1)), max(max2(s2), max2(s3))), + max(max(max2(s4), max2(s5)), max(max2(s6), max2(s7))), + ); + let old_m = t_m[row]; + let old_d = t_d[row]; + let new_m = max(old_m, tile_max); + if (row_valid) { + t_alpha[row] = exp(old_m - new_m); + let tile_sum = + exp_sum(s0, new_m) + exp_sum(s1, new_m) + + exp_sum(s2, new_m) + exp_sum(s3, new_m) + + exp_sum(s4, new_m) + exp_sum(s5, new_m) + + exp_sum(s6, new_m) + exp_sum(s7, new_m); + t_d[row] = old_d * t_alpha[row] + tile_sum; + t_m[row] = new_m; + } else { + t_alpha[row] = 0.0; + t_d[row] = 1.0; + t_m[row] = 0.0; + } + } + + load_slot = 0u; + loop { + if (load_slot >= 2u) { + break; + } + let tile_index = local_linear + load_slot * 256u; + let key_in_tile = tile_index / D4; + let load_dim4 = tile_index % D4; + let key = key_tile_start + key_in_tile; + if (key < params.context_len) { + let cache_index = + key * params.kv_token_stride4 + + kv_head * params.kv_head_stride4 + + load_dim4; + t_kv_tile[tile_index] = t_v_cache[cache_index]; + } else { + t_kv_tile[tile_index] = vec4(0.0h); + } + load_slot += 1u; + } + workgroupBarrier(); + + let alpha = t_alpha[row]; + output_acc[0] = output_acc[0] * alpha; + output_acc[1] = output_acc[1] * alpha; + output_acc[2] = output_acc[2] * alpha; + output_acc[3] = output_acc[3] * alpha; + let new_m = t_m[row]; + let row_score_base = row * 8u; + var score_block = 0u; + loop { + if (score_block >= 8u) { + break; + } + let probabilities = + exp(t_scores[row_score_base + score_block] - vec2(new_m)); + let value_key_base = score_block * 2u; + let value_dim0 = dim_vec4_base; + let value_dim1 = dim_vec4_base + 1u; + let value_dim2 = dim_vec4_base + 2u; + let value_dim3 = dim_vec4_base + 3u; + output_acc[0] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim0]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim0]) * probabilities.y; + output_acc[1] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim1]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim1]) * probabilities.y; + output_acc[2] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim2]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim2]) * probabilities.y; + output_acc[3] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim3]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim3]) * probabilities.y; + score_block += 1u; + } + workgroupBarrier(); + key_tile_start += K_TILE; + } + + if (row_valid) { + let denominator = t_d[row]; + let output_base = + token * params.o_token_stride4 + + q_head * params.o_head_stride4 + + dim_vec4_base; + t_out[output_base] = output_acc[0] / denominator; + t_out[output_base + 1u] = output_acc[1] / denominator; + t_out[output_base + 2u] = output_acc[2] / denominator; + t_out[output_base + 3u] = output_acc[3] / denominator; + } +} diff --git a/backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_q32_k16_causal_bound_wgsl.h b/backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_q32_k16_causal_bound_wgsl.h new file mode 100644 index 00000000000..e58a9d94bac --- /dev/null +++ b/backends/webgpu/runtime/ops/sdpa/streaming_attention_qwen3_q32_k16_causal_bound_wgsl.h @@ -0,0 +1,311 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from streaming_attention_qwen3_q32_k16_causal_bound.wgsl +// DO NOT EDIT. +// wgsl-sha256: 6d4396945b86cc3445698b1182fc4535ca786ad509ba128a6b6b08977e6a55e0 +inline constexpr const char* kStreamingAttentionQwen3Q32K16CausalBoundWGSL = R"( +enable f16; + +@group(0) @binding(0) var t_out: array>; +@group(0) @binding(1) var t_q: array>; +@group(0) @binding(2) var t_k_cache: array>; +@group(0) @binding(3) var t_v_cache: array>; + +struct Params { + S: u32, + context_len: u32, + input_pos: u32, + q_token_stride4: u32, + q_head_stride4: u32, + kv_token_stride4: u32, + kv_head_stride4: u32, + o_token_stride4: u32, + o_head_stride4: u32, + _pad0: u32, + _pad1: u32, + _pad2: u32, +} +@group(0) @binding(4) var params: Params; + +const HQ: u32 = 16u; +const HKV: u32 = 8u; +const G: u32 = 2u; +const D: u32 = 128u; +const D4: u32 = 32u; +const Q_TILE: u32 = 32u; +const K_TILE: u32 = 16u; +const SCALE: f32 = 0.08838834764831845; +const NEG_INF: f32 = -1.0e30; + +var t_q_tile: array, 1024>; +var t_kv_tile: array, 512>; +var t_scores: array, 256>; +var t_m: array; +var t_d: array; +var t_alpha: array; + +fn dot_qk(row: u32, key: u32) -> f32 { + let q_base = row * D4; + let k_base = key * D4; + var sum = 0.0; + var d4 = 0u; + loop { + if (d4 >= D4) { + break; + } + sum += dot(t_q_tile[q_base + d4], vec4(t_kv_tile[k_base + d4])); + d4 += 1u; + } + return sum * SCALE; +} + +fn score_for( + row: u32, + key_in_tile: u32, + key: u32, + row_valid: bool, + key_valid: bool, + token: u32, +) -> f32 { + if (row_valid && key_valid && key <= params.input_pos + token) { + return dot_qk(row, key_in_tile); + } + return NEG_INF; +} + +fn max2(v: vec2) -> f32 { + return max(v.x, v.y); +} + +fn exp_sum(v: vec2, maximum: f32) -> f32 { + let p = exp(v - vec2(maximum)); + return p.x + p.y; +} + +@compute @workgroup_size(32, 8, 1) +fn main( + @builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3, +) { + let groups_per_kv: u32 = (params.S * G + 31u) / 32u; + let kv_head = wid.x / groups_per_kv; + let row_group = wid.x % groups_per_kv; + if (kv_head >= HKV) { + return; + } + + let row: u32 = lid.x; + let dim_vec4_base: u32 = lid.y * 4u; + let logical_row: u32 = row_group * Q_TILE + row; + let row_valid = logical_row < params.S * G; + let token: u32 = logical_row / G; + let q_head: u32 = kv_head * G + logical_row % G; + let local_linear = lid.y * 32u + lid.x; + let group_max_logical_row = + min(params.S * G - 1u, row_group * Q_TILE + Q_TILE - 1u); + let group_max_token = group_max_logical_row / G; + let group_context_len = + min(params.context_len, params.input_pos + group_max_token + 1u); + + var load_slot = 0u; + loop { + if (load_slot >= 4u) { + break; + } + let tile_index = local_linear + load_slot * 256u; + let load_row = tile_index / D4; + let load_dim4 = tile_index % D4; + let load_logical_row = row_group * Q_TILE + load_row; + if (load_logical_row < params.S * G) { + let load_token = load_logical_row / G; + let load_q_head = kv_head * G + load_logical_row % G; + let q_index = + load_token * params.q_token_stride4 + + load_q_head * params.q_head_stride4 + + load_dim4; + t_q_tile[tile_index] = t_q[q_index]; + } else { + t_q_tile[tile_index] = vec4(0.0); + } + load_slot += 1u; + } + + if (lid.y == 0u) { + t_m[row] = NEG_INF; + t_d[row] = 0.0; + t_alpha[row] = 0.0; + } + workgroupBarrier(); + + var score_acc = vec2(0.0); + var output_acc: array, 4>; + output_acc[0] = vec4(0.0); + output_acc[1] = vec4(0.0); + output_acc[2] = vec4(0.0); + output_acc[3] = vec4(0.0); + + var key_tile_start = 0u; + loop { + if (key_tile_start >= group_context_len) { + break; + } + + load_slot = 0u; + loop { + if (load_slot >= 2u) { + break; + } + let tile_index = local_linear + load_slot * 256u; + let key_in_tile = tile_index / D4; + let load_dim4 = tile_index % D4; + let key = key_tile_start + key_in_tile; + if (key < params.context_len) { + let cache_index = + key * params.kv_token_stride4 + + kv_head * params.kv_head_stride4 + + load_dim4; + t_kv_tile[tile_index] = t_k_cache[cache_index]; + } else { + t_kv_tile[tile_index] = vec4(0.0h); + } + load_slot += 1u; + } + workgroupBarrier(); + + let score_key_base = lid.y * 2u; + let key0 = key_tile_start + score_key_base; + let key1 = key0 + 1u; + score_acc = vec2( + score_for(row, score_key_base, key0, row_valid, key0 < params.context_len, token), + score_for(row, score_key_base + 1u, key1, row_valid, key1 < params.context_len, token), + ); + let score_store = row * 8u + lid.y; + t_scores[score_store] = score_acc; + workgroupBarrier(); + + if (lid.y == 0u) { + let row_score_base = row * 8u; + let s0 = t_scores[row_score_base]; + let s1 = t_scores[row_score_base + 1u]; + let s2 = t_scores[row_score_base + 2u]; + let s3 = t_scores[row_score_base + 3u]; + let s4 = t_scores[row_score_base + 4u]; + let s5 = t_scores[row_score_base + 5u]; + let s6 = t_scores[row_score_base + 6u]; + let s7 = t_scores[row_score_base + 7u]; + let tile_max = max( + max(max(max2(s0), max2(s1)), max(max2(s2), max2(s3))), + max(max(max2(s4), max2(s5)), max(max2(s6), max2(s7))), + ); + let old_m = t_m[row]; + let old_d = t_d[row]; + let new_m = max(old_m, tile_max); + if (row_valid) { + t_alpha[row] = exp(old_m - new_m); + let tile_sum = + exp_sum(s0, new_m) + exp_sum(s1, new_m) + + exp_sum(s2, new_m) + exp_sum(s3, new_m) + + exp_sum(s4, new_m) + exp_sum(s5, new_m) + + exp_sum(s6, new_m) + exp_sum(s7, new_m); + t_d[row] = old_d * t_alpha[row] + tile_sum; + t_m[row] = new_m; + } else { + t_alpha[row] = 0.0; + t_d[row] = 1.0; + t_m[row] = 0.0; + } + } + + load_slot = 0u; + loop { + if (load_slot >= 2u) { + break; + } + let tile_index = local_linear + load_slot * 256u; + let key_in_tile = tile_index / D4; + let load_dim4 = tile_index % D4; + let key = key_tile_start + key_in_tile; + if (key < params.context_len) { + let cache_index = + key * params.kv_token_stride4 + + kv_head * params.kv_head_stride4 + + load_dim4; + t_kv_tile[tile_index] = t_v_cache[cache_index]; + } else { + t_kv_tile[tile_index] = vec4(0.0h); + } + load_slot += 1u; + } + workgroupBarrier(); + + let alpha = t_alpha[row]; + output_acc[0] = output_acc[0] * alpha; + output_acc[1] = output_acc[1] * alpha; + output_acc[2] = output_acc[2] * alpha; + output_acc[3] = output_acc[3] * alpha; + let new_m = t_m[row]; + let row_score_base = row * 8u; + var score_block = 0u; + loop { + if (score_block >= 8u) { + break; + } + let probabilities = + exp(t_scores[row_score_base + score_block] - vec2(new_m)); + let value_key_base = score_block * 2u; + let value_dim0 = dim_vec4_base; + let value_dim1 = dim_vec4_base + 1u; + let value_dim2 = dim_vec4_base + 2u; + let value_dim3 = dim_vec4_base + 3u; + output_acc[0] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim0]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim0]) * probabilities.y; + output_acc[1] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim1]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim1]) * probabilities.y; + output_acc[2] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim2]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim2]) * probabilities.y; + output_acc[3] += + vec4(t_kv_tile[(value_key_base + 0u) * D4 + value_dim3]) * probabilities.x + + vec4(t_kv_tile[(value_key_base + 1u) * D4 + value_dim3]) * probabilities.y; + score_block += 1u; + } + workgroupBarrier(); + key_tile_start += K_TILE; + } + + if (row_valid) { + let denominator = t_d[row]; + let output_base = + token * params.o_token_stride4 + + q_head * params.o_head_stride4 + + dim_vec4_base; + t_out[output_base] = output_acc[0] / denominator; + t_out[output_base + 1u] = output_acc[1] / denominator; + t_out[output_base + 2u] = output_acc[2] / denominator; + t_out[output_base + 3u] = output_acc[3] / denominator; + } +} +)"; + +inline constexpr uint32_t + kStreamingAttentionQwen3Q32K16CausalBoundWorkgroupSizeX = 32; +inline constexpr uint32_t + kStreamingAttentionQwen3Q32K16CausalBoundWorkgroupSizeY = 8; +inline constexpr uint32_t + kStreamingAttentionQwen3Q32K16CausalBoundWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.cpp b/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.cpp index 82ead71b597..b39c2ae9014 100644 --- a/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.cpp +++ b/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.cpp @@ -17,6 +17,7 @@ #include +#include #include #include #include @@ -70,6 +71,7 @@ void build_dispatch( WGPUBuffer uniform_buffer, uint64_t uniform_size, uint32_t workgroup_count_x, + bool retain_uniform, const char* kernel_name) { WGPUDevice device = graph.device(); @@ -101,26 +103,55 @@ void build_dispatch( graph.add_dispatch( {bundle.pipeline, bundle.bind_group, workgroup_count_x, kernel_name}); - wgpuBufferRelease(uniform_buffer); + if (retain_uniform) { + graph.own_uniform_buffer(uniform_buffer); + } else { + wgpuBufferRelease(uniform_buffer); + } +} + +FdSplitParams make_split_params(const SdpaFdDecodeState& state) { + FdSplitParams params = {}; + params.Hkv = state.Hkv; + params.D = state.D; + params.context_len = state.context_len; + params.g = state.g; + params.num_splits = state.num_splits; + params.split_len = state.split_len; + params.scale = state.scale; + return params; +} + +FdReduceParams make_reduce_params(const SdpaFdDecodeState& state) { + FdReduceParams params = {}; + params.D = state.D; + params.num_splits = state.num_splits; + return params; } } // namespace -void sdpa_fd_decode_dispatch( - WebGPUGraph& graph, - const WebGPUTensor& q, - const WebGPUTensor& k_cache, - const WebGPUTensor& v_cache, - const WebGPUTensor& out, +SdpaFdDecodeState make_sdpa_fd_decode_state( + WGPUDevice device, int64_t Hq, int64_t Hkv, int64_t D, int64_t context_len, int64_t g, float scale) { - // Defensive contract guard: the Sdpa.cpp gate only routes D <= this here, but - // keep the check (lane-owns-D reach) so a future caller can't silently - // overrun. + if (Hq <= 0 || Hkv <= 0 || D <= 0 || context_len <= 0 || g <= 0) { + throw std::runtime_error( + "WebGPU sdpa FlashDecoding: dimensions must be positive"); + } + if (Hq != Hkv * g) { + throw std::runtime_error( + "WebGPU sdpa FlashDecoding: inconsistent GQA dimensions"); + } + if (Hq > UINT32_MAX || Hkv > UINT32_MAX || D > UINT32_MAX || + context_len > UINT32_MAX || g > UINT32_MAX) { + throw std::runtime_error( + "WebGPU sdpa FlashDecoding: parameter exceeds uint32 max"); + } if (D > kSdpaFdMaxHeadDim) { throw std::runtime_error( "WebGPU sdpa FlashDecoding: head dim must be <= " + @@ -130,26 +161,59 @@ void sdpa_fd_decode_dispatch( throw std::runtime_error( "WebGPU sdpa FlashDecoding: head dim must be a multiple of 4"); } - // context_len 0 -> split_len 0 -> empty KV loop -> silent zero output; the - // Sdpa.cpp gate guarantees ctx >= 1, but fail loud if called directly. - if (context_len <= 0) { - throw std::runtime_error( - "WebGPU sdpa FlashDecoding: context_len must be positive"); - } - // Split factor: one split per kSdpaFdSplitTile KV rows, capped. uint32_t num_splits = static_cast( (context_len + kSdpaFdSplitTile - 1) / kSdpaFdSplitTile); - if (num_splits > kSdpaFdMaxSplits) { - num_splits = kSdpaFdMaxSplits; - } + num_splits = std::min(num_splits, kSdpaFdMaxSplits); const uint32_t split_len = static_cast((context_len + num_splits - 1) / num_splits); + const uint64_t split_threads = static_cast(Hq) * + static_cast(num_splits) * + static_cast(kSdpaFdSplitWorkgroupSizeX); + const uint64_t reduce_threads = + static_cast(Hq) * kSdpaFdReduceWorkgroupSizeX; + if (split_threads > UINT32_MAX || reduce_threads > UINT32_MAX) { + throw std::runtime_error( + "WebGPU sdpa FlashDecoding: thread count exceeds uint32 max"); + } + + const uint32_t split_wgc = utils::compute_1d_workgroup_count( + device, + static_cast(split_threads), + kSdpaFdSplitWorkgroupSizeX, + "fd_split"); + const uint32_t reduce_wgc = utils::compute_1d_workgroup_count( + device, + static_cast(reduce_threads), + kSdpaFdReduceWorkgroupSizeX, + "fd_reduce"); + return { + static_cast(Hq), + static_cast(Hkv), + static_cast(D), + static_cast(context_len), + static_cast(g), + num_splits, + split_len, + scale, + {split_wgc, 1u}, + {reduce_wgc, 1u}}; +} + +SdpaFdDecodeResources record_sdpa_fd_decode_dispatches( + WebGPUGraph& graph, + const WebGPUTensor& q, + const WebGPUTensor& k_cache, + const WebGPUTensor& v_cache, + const WebGPUTensor& out, + const SdpaFdDecodeState& state) { + const size_t dispatch_begin = graph.num_dispatches(); + // Scratch: per-(head,split) partials at kSdpaFdMaxSplits stride. - const uint64_t po_floats = static_cast(Hq) * - static_cast(kSdpaFdMaxSplits) * static_cast(D); - const uint64_t pml_floats = static_cast(Hq) * + const uint64_t po_floats = static_cast(state.Hq) * + static_cast(kSdpaFdMaxSplits) * static_cast(state.D); + const uint64_t pml_floats = static_cast(state.Hq) * static_cast(kSdpaFdMaxSplits) * 2ull; WGPUBuffer part_o = graph.acquire_scratch(po_floats * sizeof(float)); WebGPUGraph::ScopedScratch part_o_guard(&graph, part_o); @@ -157,14 +221,7 @@ void sdpa_fd_decode_dispatch( WebGPUGraph::ScopedScratch part_ml_guard(&graph, part_ml); // Pass 1: split (Hq*num_splits WGs) -> writes part_o, part_ml. - FdSplitParams sp = {}; - sp.Hkv = static_cast(Hkv); - sp.D = static_cast(D); - sp.context_len = static_cast(context_len); - sp.g = static_cast(g); - sp.num_splits = num_splits; - sp.split_len = split_len; - sp.scale = scale; + FdSplitParams sp = make_split_params(state); WGPUBuffer ub_split = graph.make_uniform_buffer(&sp, sizeof(sp)); BufferBinding split_bindings[5] = { {part_o, po_floats * sizeof(float)}, @@ -172,20 +229,6 @@ void sdpa_fd_decode_dispatch( {q.buffer, q.nbytes}, {k_cache.buffer, k_cache.nbytes}, {v_cache.buffer, v_cache.nbytes}}; - // Compute the thread product in 64-bit + guard before the u32 cast, mirroring - // the Sdpa.cpp aw_floats > UINT32_MAX guards. - const uint64_t split_threads = static_cast(Hq) * - static_cast(num_splits) * - static_cast(kSdpaFdSplitWorkgroupSizeX); - if (split_threads > UINT32_MAX) { - throw std::runtime_error( - "WebGPU sdpa FlashDecoding: split thread count exceeds uint32 max"); - } - const uint32_t wgc_split = utils::compute_1d_workgroup_count( - graph.device(), - static_cast(split_threads), - kSdpaFdSplitWorkgroupSizeX, - "fd_split"); const char* split_shader = kSdpaFdSplitWGSL; if (graph.kv_f16()) { split_shader = kSdpaFdSplitHalfWGSL; @@ -198,23 +241,17 @@ void sdpa_fd_decode_dispatch( 2, ub_split, sizeof(sp), - wgc_split, + state.split_grid.x, + true, "fd_split"); // Pass 2: reduce (Hq WGs) -> reads part_o, part_ml; writes out. - FdReduceParams rp = {}; - rp.D = static_cast(D); - rp.num_splits = num_splits; + FdReduceParams rp = make_reduce_params(state); WGPUBuffer ub_reduce = graph.make_uniform_buffer(&rp, sizeof(rp)); BufferBinding reduce_bindings[3] = { {out.buffer, out.nbytes}, {part_o, po_floats * sizeof(float)}, {part_ml, pml_floats * sizeof(float)}}; - const uint32_t wgc_reduce = utils::compute_1d_workgroup_count( - graph.device(), - static_cast(Hq) * kSdpaFdReduceWorkgroupSizeX, - kSdpaFdReduceWorkgroupSizeX, - "fd_reduce"); build_dispatch( graph, kSdpaFdReduceWGSL, @@ -223,8 +260,27 @@ void sdpa_fd_decode_dispatch( 1, ub_reduce, sizeof(rp), - wgc_reduce, + state.reduce_grid.x, + true, "fd_reduce"); + + return {ub_split, ub_reduce, {dispatch_begin, graph.num_dispatches()}}; +} + +void write_sdpa_fd_decode_uniforms( + WGPUQueue queue, + const SdpaFdDecodeResources& resources, + const SdpaFdDecodeState& state) { + FdSplitParams split_params = make_split_params(state); + FdReduceParams reduce_params = make_reduce_params(state); + wgpuQueueWriteBuffer( + queue, resources.split_uniform, 0, &split_params, sizeof(split_params)); + wgpuQueueWriteBuffer( + queue, + resources.reduce_uniform, + 0, + &reduce_params, + sizeof(reduce_params)); } } // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.h b/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.h index 1b161cd8ec0..d86b4c624b8 100644 --- a/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.h +++ b/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.h @@ -9,6 +9,7 @@ #pragma once #include +#include #include @@ -19,15 +20,27 @@ namespace executorch::backends::webgpu { // SDPA path (the FD selection predicate in Sdpa.cpp checks this). constexpr int64_t kSdpaFdMaxHeadDim = 128; -// Split-KV FlashDecoding decode dispatch (S==1): a split pass over -// Hq*num_splits workgroups + a reduce pass over Hq workgroups. Called from the -// Sdpa.cpp WEBGPU_SDPA_FD branch. -void sdpa_fd_decode_dispatch( - WebGPUGraph& graph, - const WebGPUTensor& q, - const WebGPUTensor& k_cache, - const WebGPUTensor& v_cache, - const WebGPUTensor& out, +struct SdpaFdDecodeState { + uint32_t Hq; + uint32_t Hkv; + uint32_t D; + uint32_t context_len; + uint32_t g; + uint32_t num_splits; + uint32_t split_len; + float scale; + utils::WgCount split_grid; + utils::WgCount reduce_grid; +}; + +struct SdpaFdDecodeResources { + WGPUBuffer split_uniform; + WGPUBuffer reduce_uniform; + utils::DispatchRange dispatch_range; +}; + +SdpaFdDecodeState make_sdpa_fd_decode_state( + WGPUDevice device, int64_t Hq, int64_t Hkv, int64_t D, @@ -35,4 +48,19 @@ void sdpa_fd_decode_dispatch( int64_t g, float scale); +// Records split + reduce with retained UBOs. Route selection is owned by the +// caller so this helper never mutates recorded dispatch counts. +SdpaFdDecodeResources record_sdpa_fd_decode_dispatches( + WebGPUGraph& graph, + const WebGPUTensor& q, + const WebGPUTensor& k_cache, + const WebGPUTensor& v_cache, + const WebGPUTensor& out, + const SdpaFdDecodeState& state); + +void write_sdpa_fd_decode_uniforms( + WGPUQueue queue, + const SdpaFdDecodeResources& resources, + const SdpaFdDecodeState& state); + } // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/to_copy/ToCopy.cpp b/backends/webgpu/runtime/ops/to_copy/ToCopy.cpp index 05f7ff6fe24..3f9c8e29e3c 100644 --- a/backends/webgpu/runtime/ops/to_copy/ToCopy.cpp +++ b/backends/webgpu/runtime/ops/to_copy/ToCopy.cpp @@ -10,12 +10,14 @@ #include #include #include +#include #include #include #include #include +#include #include #include #include @@ -118,6 +120,98 @@ void add_convert_op( graph.own_uniform_buffer(uniform_buffer); } +// Decode byte-packed bool storage into numeric fp32 values. +void add_bool_to_float_op(WebGPUGraph& graph, int in_id, int out_id) { + WGPUDevice device = graph.device(); + const auto& in_tensor = graph.get_tensor(in_id); + const auto& out_tensor = graph.get_tensor(out_id); + if (in_tensor.buffer == nullptr || out_tensor.buffer == nullptr) { + throw std::runtime_error("to_copy_bool_to_float: null buffer binding"); + } + if (!in_tensor.is_bool || in_tensor.elem_size != 1 || out_tensor.is_int || + out_tensor.elem_size != sizeof(float) || + out_tensor.nbytes % sizeof(float) != 0 || + out_tensor.nbytes / sizeof(float) != in_tensor.nbytes) { + throw std::runtime_error("to_copy_bool_to_float: dtype/numel mismatch"); + } + if (in_tensor.nbytes == 0u || + in_tensor.nbytes > std::numeric_limits::max()) { + throw std::runtime_error( + "to_copy_bool_to_float: numel must be nonzero and fit u32"); + } + + const uint32_t num_elements = static_cast(in_tensor.nbytes); + const uint64_t input_bind_size_u64 = + (static_cast(in_tensor.nbytes) + 3u) & ~uint64_t(3); + if (input_bind_size_u64 > std::numeric_limits::max()) { + throw std::runtime_error( + "to_copy_bool_to_float: input binding size overflows"); + } + const size_t input_bind_size = static_cast(input_bind_size_u64); + + const uint32_t wg_size = + utils::clamp_workgroup_size(device, kToCopyBoolToFloatWorkgroupSizeX); + const uint32_t workgroup_count = utils::compute_1d_workgroup_count( + device, num_elements, wg_size, "to_copy_bool_to_float"); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + ConvertParams params = {}; + params.num_elements = num_elements; + WGPUBuffer uniform_buffer = + utils::make_uniform(device, ¶ms, sizeof(ConvertParams)); + graph.add_uniform_buffer_bytes(sizeof(ConvertParams)); + + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kToCopyBoolToFloatWGSL, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + in_tensor.buffer, + input_bind_size}, + {1, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {2, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(ConvertParams)}, + }, + &wg_size_constant, + 1); + + const size_t dispatch_idx = + graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count}); + + WGPUBuffer params_buf = uniform_buffer; + graph.add_tensor_resize_hook( + in_id, + [in_id, out_id, wg_size, dispatch_idx, params_buf](WebGPUGraph& g) { + const auto& dims = g.cur_dims(in_id); + const uint64_t numel = utils::numel_of(dims); + if (numel == 0u || numel > std::numeric_limits::max()) { + throw std::runtime_error( + "to_copy_bool_to_float(resize): invalid numel"); + } + g.set_cur_dims(out_id, dims); + ConvertParams p = {}; + p.num_elements = static_cast(numel); + wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); + g.dispatch_at(dispatch_idx).workgroup_count_x = + utils::compute_1d_workgroup_count( + g.device(), + static_cast(numel), + wg_size, + "to_copy_bool_to_float(resize)"); + }); + + graph.own_uniform_buffer(uniform_buffer); +} + void to_copy_impl(WebGPUGraph& graph, const std::vector& args) { // aten._to_copy.default args: [self, ...kwargs, out]; out = last value id. add_to_copy_node(graph, args.at(0), args.at(args.size() - 1)); @@ -137,7 +231,10 @@ void add_to_copy_node(WebGPUGraph& graph, int in_id, int out_id) { } // int<->float = numeric convert (mirrors Vulkan add_view_copy_convert_node). - if (in_tensor.is_int && !out_tensor.is_int) { + if (in_tensor.is_bool && !out_tensor.is_int && + out_tensor.elem_size == sizeof(float)) { + add_bool_to_float_op(graph, in_id, out_id); + } else if (in_tensor.is_int && !out_tensor.is_int) { add_convert_op( graph, in_id, diff --git a/backends/webgpu/runtime/ops/to_copy/to_copy_bool_to_float.wgsl b/backends/webgpu/runtime/ops/to_copy/to_copy_bool_to_float.wgsl new file mode 100644 index 00000000000..239730de65d --- /dev/null +++ b/backends/webgpu/runtime/ops/to_copy/to_copy_bool_to_float.wgsl @@ -0,0 +1,24 @@ +override wg_size: u32 = 256u; + +struct Params { + num_elements: u32, + _pad0: u32, + _pad1: u32, + _pad2: u32, +}; + +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; +@group(0) @binding(2) var params: Params; + +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let idx = gid.x; + if (idx >= params.num_elements) { + return; + } + let word = input[idx / 4u]; + let byte_shift = (idx % 4u) * 8u; + let value = (word >> byte_shift) & 0xffu; + output[idx] = select(0.0, 1.0, value != 0u); +} diff --git a/backends/webgpu/runtime/ops/to_copy/to_copy_bool_to_float_wgsl.h b/backends/webgpu/runtime/ops/to_copy/to_copy_bool_to_float_wgsl.h new file mode 100644 index 00000000000..ef7e40976f3 --- /dev/null +++ b/backends/webgpu/runtime/ops/to_copy/to_copy_bool_to_float_wgsl.h @@ -0,0 +1,48 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from to_copy_bool_to_float.wgsl - DO NOT EDIT. +// wgsl-sha256: 29fd43b2f638489e9b8d72b2cc9140d07174c750cbdc291e63793319a2fa5961 +inline constexpr const char* kToCopyBoolToFloatWGSL = R"( +override wg_size: u32 = 256u; + +struct Params { + num_elements: u32, + _pad0: u32, + _pad1: u32, + _pad2: u32, +}; + +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; +@group(0) @binding(2) var params: Params; + +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let idx = gid.x; + if (idx >= params.num_elements) { + return; + } + let word = input[idx / 4u]; + let byte_shift = (idx % 4u) * 8u; + let value = (word >> byte_shift) & 0xffu; + output[idx] = select(0.0, 1.0, value != 0u); +} +)"; + +inline constexpr uint32_t kToCopyBoolToFloatWorkgroupSizeX = 256; +inline constexpr uint32_t kToCopyBoolToFloatWorkgroupSizeY = 1; +inline constexpr uint32_t kToCopyBoolToFloatWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/to_copy/to_copy_float_to_int.wgsl b/backends/webgpu/runtime/ops/to_copy/to_copy_convert.wgsl similarity index 62% rename from backends/webgpu/runtime/ops/to_copy/to_copy_float_to_int.wgsl rename to backends/webgpu/runtime/ops/to_copy/to_copy_convert.wgsl index 3eb6cb44595..f1113f0e14c 100644 --- a/backends/webgpu/runtime/ops/to_copy/to_copy_float_to_int.wgsl +++ b/backends/webgpu/runtime/ops/to_copy/to_copy_convert.wgsl @@ -1,5 +1,5 @@ -@group(0) @binding(0) var input: array; -@group(0) @binding(1) var output: array; +@group(0) @binding(0) var input: array<${IN_TYPE}>; +@group(0) @binding(1) var output: array<${OUT_TYPE}>; struct Params { num_elements: u32, @@ -14,5 +14,5 @@ fn main(@builtin(global_invocation_id) gid: vec3) { if (idx >= params.num_elements) { return; } - output[idx] = i32(input[idx]); + output[idx] = ${OUT_TYPE}(input[idx]); } diff --git a/backends/webgpu/runtime/ops/to_copy/to_copy_convert.yaml b/backends/webgpu/runtime/ops/to_copy/to_copy_convert.yaml new file mode 100644 index 00000000000..219b2d8edd3 --- /dev/null +++ b/backends/webgpu/runtime/ops/to_copy/to_copy_convert.yaml @@ -0,0 +1,15 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +to_copy_convert: + parameter_names_with_default_values: + IN_TYPE: f32 + OUT_TYPE: i32 + shader_variants: + - NAME: to_copy_float_to_int + - NAME: to_copy_int_to_float + IN_TYPE: i32 + OUT_TYPE: f32 diff --git a/backends/webgpu/runtime/ops/to_copy/to_copy_float_to_int_wgsl.h b/backends/webgpu/runtime/ops/to_copy/to_copy_float_to_int_wgsl.h index e7e0391dd13..1a384c747d8 100644 --- a/backends/webgpu/runtime/ops/to_copy/to_copy_float_to_int_wgsl.h +++ b/backends/webgpu/runtime/ops/to_copy/to_copy_float_to_int_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from to_copy_float_to_int.wgsl - DO NOT EDIT. +// @generated from to_copy_convert.wgsl - DO NOT EDIT. // wgsl-sha256: c331e00e3171eecbe6317ac9df0a5f9cd6d25da26a9a587250f1cc6086dc3c8f inline constexpr const char* kToCopyFloatToIntWGSL = R"( @group(0) @binding(0) var input: array; diff --git a/backends/webgpu/runtime/ops/to_copy/to_copy_int_to_float.wgsl b/backends/webgpu/runtime/ops/to_copy/to_copy_int_to_float.wgsl deleted file mode 100644 index 87affe78290..00000000000 --- a/backends/webgpu/runtime/ops/to_copy/to_copy_int_to_float.wgsl +++ /dev/null @@ -1,18 +0,0 @@ -@group(0) @binding(0) var input: array; -@group(0) @binding(1) var output: array; - -struct Params { - num_elements: u32, -} -@group(0) @binding(2) var params: Params; - -override wg_size: u32 = 64u; - -@compute @workgroup_size(wg_size, 1, 1) -fn main(@builtin(global_invocation_id) gid: vec3) { - let idx = gid.x; - if (idx >= params.num_elements) { - return; - } - output[idx] = f32(input[idx]); -} diff --git a/backends/webgpu/runtime/ops/to_copy/to_copy_int_to_float_wgsl.h b/backends/webgpu/runtime/ops/to_copy/to_copy_int_to_float_wgsl.h index fd18700f17c..6fdf37ec2b7 100644 --- a/backends/webgpu/runtime/ops/to_copy/to_copy_int_to_float_wgsl.h +++ b/backends/webgpu/runtime/ops/to_copy/to_copy_int_to_float_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from to_copy_int_to_float.wgsl - DO NOT EDIT. +// @generated from to_copy_convert.wgsl - DO NOT EDIT. // wgsl-sha256: e18dd733a3838f83eded4977a2a2b21119099c8409b234f12474fae5acc9b195 inline constexpr const char* kToCopyIntToFloatWGSL = R"( @group(0) @binding(0) var input: array; diff --git a/backends/webgpu/runtime/ops/unary/abs.wgsl b/backends/webgpu/runtime/ops/unary/abs.wgsl deleted file mode 100644 index e3e10c75dd9..00000000000 --- a/backends/webgpu/runtime/ops/unary/abs.wgsl +++ /dev/null @@ -1,21 +0,0 @@ -@group(0) @binding(0) var input: array; -@group(0) @binding(1) var output: array; - -struct Params { - num_elements: u32, -} -@group(0) @binding(2) var params: Params; - -override wg_size: u32 = 256u; - -@compute @workgroup_size(wg_size) -fn main( - @builtin(global_invocation_id) gid: vec3, - @builtin(num_workgroups) num_workgroups: vec3) { - let idx = gid.x + gid.y * (num_workgroups.x * wg_size); - if (idx >= params.num_elements) { - return; - } - let x = input[idx]; - output[idx] = abs(x); -} diff --git a/backends/webgpu/runtime/ops/unary/abs_wgsl.h b/backends/webgpu/runtime/ops/unary/abs_wgsl.h index 3d2873e69c4..0efb72d5189 100644 --- a/backends/webgpu/runtime/ops/unary/abs_wgsl.h +++ b/backends/webgpu/runtime/ops/unary/abs_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from abs.wgsl - DO NOT EDIT. +// @generated from unary.wgsl - DO NOT EDIT. // wgsl-sha256: 39d3c163fdf6a92286828f4b3217e00294e3ca5634a878ed5fd34e3b1cdf0a27 inline constexpr const char* kAbsWGSL = R"( @group(0) @binding(0) var input: array; diff --git a/backends/webgpu/runtime/ops/unary/cos_wgsl.h b/backends/webgpu/runtime/ops/unary/cos_wgsl.h index 422b5618146..4ca99df88f3 100644 --- a/backends/webgpu/runtime/ops/unary/cos_wgsl.h +++ b/backends/webgpu/runtime/ops/unary/cos_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from cos.wgsl - DO NOT EDIT. +// @generated from unary.wgsl - DO NOT EDIT. // wgsl-sha256: 9df78873e5fae98d347c26db2a02b047ea3d5d2c93f0761cb9ac6995f9a71ab2 inline constexpr const char* kCosWGSL = R"( @group(0) @binding(0) var input: array; diff --git a/backends/webgpu/runtime/ops/unary/exp.wgsl b/backends/webgpu/runtime/ops/unary/exp.wgsl deleted file mode 100644 index b69aa509d8e..00000000000 --- a/backends/webgpu/runtime/ops/unary/exp.wgsl +++ /dev/null @@ -1,21 +0,0 @@ -@group(0) @binding(0) var input: array; -@group(0) @binding(1) var output: array; - -struct Params { - num_elements: u32, -} -@group(0) @binding(2) var params: Params; - -override wg_size: u32 = 256u; - -@compute @workgroup_size(wg_size) -fn main( - @builtin(global_invocation_id) gid: vec3, - @builtin(num_workgroups) num_workgroups: vec3) { - let idx = gid.x + gid.y * (num_workgroups.x * wg_size); - if (idx >= params.num_elements) { - return; - } - let x = input[idx]; - output[idx] = exp(x); -} diff --git a/backends/webgpu/runtime/ops/unary/exp_wgsl.h b/backends/webgpu/runtime/ops/unary/exp_wgsl.h index 28b6c67cb18..cbf85fd415a 100644 --- a/backends/webgpu/runtime/ops/unary/exp_wgsl.h +++ b/backends/webgpu/runtime/ops/unary/exp_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from exp.wgsl - DO NOT EDIT. +// @generated from unary.wgsl - DO NOT EDIT. // wgsl-sha256: 3171399bc36acf9c1cb2a03c2a31038318203c4c63ab03c4881df7a660346020 inline constexpr const char* kExpWGSL = R"( @group(0) @binding(0) var input: array; diff --git a/backends/webgpu/runtime/ops/unary/hardswish.wgsl b/backends/webgpu/runtime/ops/unary/hardswish.wgsl deleted file mode 100644 index 2278caf3664..00000000000 --- a/backends/webgpu/runtime/ops/unary/hardswish.wgsl +++ /dev/null @@ -1,21 +0,0 @@ -@group(0) @binding(0) var input: array; -@group(0) @binding(1) var output: array; - -struct Params { - num_elements: u32, -} -@group(0) @binding(2) var params: Params; - -override wg_size: u32 = 256u; - -@compute @workgroup_size(wg_size) -fn main( - @builtin(global_invocation_id) gid: vec3, - @builtin(num_workgroups) num_workgroups: vec3) { - let idx = gid.x + gid.y * (num_workgroups.x * wg_size); - if (idx >= params.num_elements) { - return; - } - let x = input[idx]; - output[idx] = select(select(x * (x + 3.0) / 6.0, x, x >= 3.0), 0.0, x <= -3.0); -} diff --git a/backends/webgpu/runtime/ops/unary/hardswish_wgsl.h b/backends/webgpu/runtime/ops/unary/hardswish_wgsl.h index 0c991547b9a..43f104c7f8f 100644 --- a/backends/webgpu/runtime/ops/unary/hardswish_wgsl.h +++ b/backends/webgpu/runtime/ops/unary/hardswish_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from hardswish.wgsl - DO NOT EDIT. +// @generated from unary.wgsl - DO NOT EDIT. // wgsl-sha256: c874a15ef6cdaec71187296016cc2a1515f5e7c889b97dfa8fd4b278e6e2c3d5 inline constexpr const char* kHardswishWGSL = R"( @group(0) @binding(0) var input: array; diff --git a/backends/webgpu/runtime/ops/unary/neg.wgsl b/backends/webgpu/runtime/ops/unary/neg.wgsl deleted file mode 100644 index c977957957d..00000000000 --- a/backends/webgpu/runtime/ops/unary/neg.wgsl +++ /dev/null @@ -1,21 +0,0 @@ -@group(0) @binding(0) var input: array; -@group(0) @binding(1) var output: array; - -struct Params { - num_elements: u32, -} -@group(0) @binding(2) var params: Params; - -override wg_size: u32 = 256u; - -@compute @workgroup_size(wg_size) -fn main( - @builtin(global_invocation_id) gid: vec3, - @builtin(num_workgroups) num_workgroups: vec3) { - let idx = gid.x + gid.y * (num_workgroups.x * wg_size); - if (idx >= params.num_elements) { - return; - } - let x = input[idx]; - output[idx] = -x; -} diff --git a/backends/webgpu/runtime/ops/unary/neg_wgsl.h b/backends/webgpu/runtime/ops/unary/neg_wgsl.h index d528c45fea0..c4c7bb989ac 100644 --- a/backends/webgpu/runtime/ops/unary/neg_wgsl.h +++ b/backends/webgpu/runtime/ops/unary/neg_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from neg.wgsl - DO NOT EDIT. +// @generated from unary.wgsl - DO NOT EDIT. // wgsl-sha256: 8851b9f42d14153f6f04484fee2f8bf67bda26dea892ff48768e09e6ad49cee1 inline constexpr const char* kNegWGSL = R"( @group(0) @binding(0) var input: array; diff --git a/backends/webgpu/runtime/ops/unary/round.wgsl b/backends/webgpu/runtime/ops/unary/round.wgsl deleted file mode 100644 index 2269ca59988..00000000000 --- a/backends/webgpu/runtime/ops/unary/round.wgsl +++ /dev/null @@ -1,21 +0,0 @@ -@group(0) @binding(0) var input: array; -@group(0) @binding(1) var output: array; - -struct Params { - num_elements: u32, -} -@group(0) @binding(2) var params: Params; - -override wg_size: u32 = 256u; - -@compute @workgroup_size(wg_size) -fn main( - @builtin(global_invocation_id) gid: vec3, - @builtin(num_workgroups) num_workgroups: vec3) { - let idx = gid.x + gid.y * (num_workgroups.x * wg_size); - if (idx >= params.num_elements) { - return; - } - let x = input[idx]; - output[idx] = round(x); -} diff --git a/backends/webgpu/runtime/ops/unary/round_wgsl.h b/backends/webgpu/runtime/ops/unary/round_wgsl.h index 209305bc855..8c805ba4f80 100644 --- a/backends/webgpu/runtime/ops/unary/round_wgsl.h +++ b/backends/webgpu/runtime/ops/unary/round_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from round.wgsl - DO NOT EDIT. +// @generated from unary.wgsl - DO NOT EDIT. // wgsl-sha256: 8f3e0edbeb81aa50f35e691c78554e8057fa8d78fe8a86454f4f42e5e8871452 inline constexpr const char* kRoundWGSL = R"( @group(0) @binding(0) var input: array; diff --git a/backends/webgpu/runtime/ops/unary/rsqrt.wgsl b/backends/webgpu/runtime/ops/unary/rsqrt.wgsl deleted file mode 100644 index 1f50c4f66ac..00000000000 --- a/backends/webgpu/runtime/ops/unary/rsqrt.wgsl +++ /dev/null @@ -1,21 +0,0 @@ -@group(0) @binding(0) var input: array; -@group(0) @binding(1) var output: array; - -struct Params { - num_elements: u32, -} -@group(0) @binding(2) var params: Params; - -override wg_size: u32 = 256u; - -@compute @workgroup_size(wg_size) -fn main( - @builtin(global_invocation_id) gid: vec3, - @builtin(num_workgroups) num_workgroups: vec3) { - let idx = gid.x + gid.y * (num_workgroups.x * wg_size); - if (idx >= params.num_elements) { - return; - } - let x = input[idx]; - output[idx] = 1.0 / sqrt(x); -} diff --git a/backends/webgpu/runtime/ops/unary/rsqrt_wgsl.h b/backends/webgpu/runtime/ops/unary/rsqrt_wgsl.h index 83e74a1b8cf..58bd05d010e 100644 --- a/backends/webgpu/runtime/ops/unary/rsqrt_wgsl.h +++ b/backends/webgpu/runtime/ops/unary/rsqrt_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from rsqrt.wgsl - DO NOT EDIT. +// @generated from unary.wgsl - DO NOT EDIT. // wgsl-sha256: 108765d5a23b87473f34651875d08abf2a5fa8980bd92fc8cbe3617295097747 inline constexpr const char* kRsqrtWGSL = R"( @group(0) @binding(0) var input: array; diff --git a/backends/webgpu/runtime/ops/unary/sin.wgsl b/backends/webgpu/runtime/ops/unary/sin.wgsl deleted file mode 100644 index ffd2a07ea8a..00000000000 --- a/backends/webgpu/runtime/ops/unary/sin.wgsl +++ /dev/null @@ -1,21 +0,0 @@ -@group(0) @binding(0) var input: array; -@group(0) @binding(1) var output: array; - -struct Params { - num_elements: u32, -} -@group(0) @binding(2) var params: Params; - -override wg_size: u32 = 256u; - -@compute @workgroup_size(wg_size) -fn main( - @builtin(global_invocation_id) gid: vec3, - @builtin(num_workgroups) num_workgroups: vec3) { - let idx = gid.x + gid.y * (num_workgroups.x * wg_size); - if (idx >= params.num_elements) { - return; - } - let x = input[idx]; - output[idx] = sin(x); -} diff --git a/backends/webgpu/runtime/ops/unary/sin_wgsl.h b/backends/webgpu/runtime/ops/unary/sin_wgsl.h index f22229b2342..54184a1ccb8 100644 --- a/backends/webgpu/runtime/ops/unary/sin_wgsl.h +++ b/backends/webgpu/runtime/ops/unary/sin_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from sin.wgsl - DO NOT EDIT. +// @generated from unary.wgsl - DO NOT EDIT. // wgsl-sha256: e5762804773659d348fddddcef4935807ae6fe7d92c92eb17a2f44aae8f2c5b9 inline constexpr const char* kSinWGSL = R"( @group(0) @binding(0) var input: array; diff --git a/backends/webgpu/runtime/ops/unary/sqrt.wgsl b/backends/webgpu/runtime/ops/unary/sqrt.wgsl deleted file mode 100644 index e34ff440007..00000000000 --- a/backends/webgpu/runtime/ops/unary/sqrt.wgsl +++ /dev/null @@ -1,21 +0,0 @@ -@group(0) @binding(0) var input: array; -@group(0) @binding(1) var output: array; - -struct Params { - num_elements: u32, -} -@group(0) @binding(2) var params: Params; - -override wg_size: u32 = 256u; - -@compute @workgroup_size(wg_size) -fn main( - @builtin(global_invocation_id) gid: vec3, - @builtin(num_workgroups) num_workgroups: vec3) { - let idx = gid.x + gid.y * (num_workgroups.x * wg_size); - if (idx >= params.num_elements) { - return; - } - let x = input[idx]; - output[idx] = sqrt(x); -} diff --git a/backends/webgpu/runtime/ops/unary/sqrt_wgsl.h b/backends/webgpu/runtime/ops/unary/sqrt_wgsl.h index 260a4ce4266..42dcdb838c0 100644 --- a/backends/webgpu/runtime/ops/unary/sqrt_wgsl.h +++ b/backends/webgpu/runtime/ops/unary/sqrt_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from sqrt.wgsl - DO NOT EDIT. +// @generated from unary.wgsl - DO NOT EDIT. // wgsl-sha256: 008534ae365969f5c180b42e8d6d0b131df78f181e5435abbcafc3ffb8be8aac inline constexpr const char* kSqrtWGSL = R"( @group(0) @binding(0) var input: array; diff --git a/backends/webgpu/runtime/ops/unary/tanh.wgsl b/backends/webgpu/runtime/ops/unary/tanh.wgsl deleted file mode 100644 index 5a541e21699..00000000000 --- a/backends/webgpu/runtime/ops/unary/tanh.wgsl +++ /dev/null @@ -1,21 +0,0 @@ -@group(0) @binding(0) var input: array; -@group(0) @binding(1) var output: array; - -struct Params { - num_elements: u32, -} -@group(0) @binding(2) var params: Params; - -override wg_size: u32 = 256u; - -@compute @workgroup_size(wg_size) -fn main( - @builtin(global_invocation_id) gid: vec3, - @builtin(num_workgroups) num_workgroups: vec3) { - let idx = gid.x + gid.y * (num_workgroups.x * wg_size); - if (idx >= params.num_elements) { - return; - } - let x = input[idx]; - output[idx] = tanh(clamp(x, -15.0, 15.0)); -} diff --git a/backends/webgpu/runtime/ops/unary/tanh_wgsl.h b/backends/webgpu/runtime/ops/unary/tanh_wgsl.h index eef1ecd91af..51ba4d3919c 100644 --- a/backends/webgpu/runtime/ops/unary/tanh_wgsl.h +++ b/backends/webgpu/runtime/ops/unary/tanh_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from tanh.wgsl - DO NOT EDIT. +// @generated from unary.wgsl - DO NOT EDIT. // wgsl-sha256: 5bd7eb1c6411940d84a9b311884f35b39f15b82103b14bab02902290ed6b0339 inline constexpr const char* kTanhWGSL = R"( @group(0) @binding(0) var input: array; diff --git a/backends/webgpu/runtime/ops/unary/cos.wgsl b/backends/webgpu/runtime/ops/unary/unary.wgsl similarity index 94% rename from backends/webgpu/runtime/ops/unary/cos.wgsl rename to backends/webgpu/runtime/ops/unary/unary.wgsl index c2bafe7b248..d974a2f3319 100644 --- a/backends/webgpu/runtime/ops/unary/cos.wgsl +++ b/backends/webgpu/runtime/ops/unary/unary.wgsl @@ -17,5 +17,5 @@ fn main( return; } let x = input[idx]; - output[idx] = cos(x); + output[idx] = ${OPERATOR}; } diff --git a/backends/webgpu/runtime/ops/unary/unary.yaml b/backends/webgpu/runtime/ops/unary/unary.yaml new file mode 100644 index 00000000000..84c0dd660d7 --- /dev/null +++ b/backends/webgpu/runtime/ops/unary/unary.yaml @@ -0,0 +1,29 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +unary: + parameter_names_with_default_values: + OPERATOR: "abs(x)" + shader_variants: + - NAME: abs + - NAME: cos + OPERATOR: "cos(x)" + - NAME: exp + OPERATOR: "exp(x)" + - NAME: hardswish + OPERATOR: "select(select(x * (x + 3.0) / 6.0, x, x >= 3.0), 0.0, x <= -3.0)" + - NAME: neg + OPERATOR: "-x" + - NAME: round + OPERATOR: "round(x)" + - NAME: rsqrt + OPERATOR: "1.0 / sqrt(x)" + - NAME: sin + OPERATOR: "sin(x)" + - NAME: sqrt + OPERATOR: "sqrt(x)" + - NAME: tanh + OPERATOR: "tanh(clamp(x, -15.0, 15.0))" diff --git a/backends/webgpu/runtime/passes/QkvBk64.cpp b/backends/webgpu/runtime/passes/QkvBk64.cpp new file mode 100644 index 00000000000..73c129e45b9 --- /dev/null +++ b/backends/webgpu/runtime/passes/QkvBk64.cpp @@ -0,0 +1,444 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include + +#include +#include + +#include +#include +#include + +namespace executorch::backends::webgpu::passes { + +namespace { + +constexpr const char* kQ4gswLinearOpName = "et_vk.linear_q4gsw.default"; + +constexpr uint32_t kQkvQWidth = 2048u; +constexpr uint32_t kQkvKvWidth = 512u; +constexpr uint32_t kQkvFusedWidth = 3072u; +constexpr uint32_t kQkvK = 2048u; +constexpr uint32_t kQkvKPacked = 1024u; +constexpr uint32_t kQkvGroupSize = 64u; +constexpr uint32_t kQkvNumGroups = 32u; +constexpr uint32_t kQkvTile = 64u; + +// Uniform layout matching q4gsw_qkv_bk64.wgsl's Params struct. +struct QkvBk64Params { + uint32_t M; + uint32_t N; + uint32_t K; + uint32_t K_packed; + uint32_t group_size; + uint32_t padded_N; + uint32_t has_bias; + uint32_t _pad; +}; +static_assert(sizeof(QkvBk64Params) == 32); + +bool is_qkv_bk64_live_m(uint32_t m) { + return m == 128u || m == 508u || m == 512u; +} + +struct QkvBk64ResizeContext { + int input_id; + std::array output_ids; + std::array separate_begin; + std::array separate_end; + size_t fused_dispatch; + uint32_t max_m; + WGPUBuffer params_buffer; +}; + +void resize_qkv_bk64(WebGPUGraph& graph, const QkvBk64ResizeContext& context) { + const auto& input_dims = graph.cur_dims(context.input_id); + const uint64_t input_numel = utils::numel_of(input_dims); + if (input_dims.empty() || input_numel % kQkvK != 0u) { + throw std::runtime_error( + "linear_q4gsw_bk64_qkv(resize): malformed input shape"); + } + const uint64_t live_m = input_numel / kQkvK; + if (live_m == 0u || live_m > context.max_m) { + throw std::runtime_error( + "linear_q4gsw_bk64_qkv(resize): live M out of range"); + } + const uint32_t m = static_cast(live_m); + const uint32_t widths[3] = {kQkvQWidth, kQkvKvWidth, kQkvKvWidth}; + for (size_t i = 0; i < context.output_ids.size(); i++) { + std::vector output_dims = input_dims; + output_dims.back() = widths[i]; + graph.set_cur_dims(context.output_ids[i], output_dims); + } + + const QkvBk64Params params = { + m, + kQkvFusedWidth, + kQkvK, + kQkvKPacked, + kQkvGroupSize, + kQkvFusedWidth, + 0u, + 0u}; + wgpuQueueWriteBuffer( + graph.queue(), context.params_buffer, 0, ¶ms, sizeof(params)); + + const bool use_fused = is_qkv_bk64_live_m(m); + auto& fused = graph.dispatch_at(context.fused_dispatch); + fused.workgroup_count_x = use_fused + ? ((m + kQkvTile - 1u) / kQkvTile) * (kQkvFusedWidth / kQkvTile) + : 0u; + fused.workgroup_count_y = use_fused ? 1u : 0u; + if (use_fused) { + // The separate projections are inactive while the fused route is live. + for (size_t member = 0; member < context.separate_begin.size(); member++) { + for (size_t i = context.separate_begin[member]; + i < context.separate_end[member]; + i++) { + auto& dispatch = graph.dispatch_at(i); + dispatch.workgroup_count_x = 0u; + dispatch.workgroup_count_y = 0u; + } + } + } else { + // The separate projections' own resize hooks are registered before this + // one (Phase 3 processes each Q/K/V member before this fusion's combined + // hook) and unconditionally restore their live grids, so this hook + // normally has nothing to do here. That ordering isn't enforced by the + // type system, so fail loud rather than silently drop Q/K/V outputs if a + // future change ever violates it. + for (size_t member = 0; member < context.separate_begin.size(); member++) { + for (size_t i = context.separate_begin[member]; + i < context.separate_end[member]; + i++) { + const auto& dispatch = graph.dispatch_at(i); + if (dispatch.workgroup_count_x == 0u || + dispatch.workgroup_count_y == 0u) { + throw std::runtime_error( + "linear_q4gsw_bk64_qkv(resize): separate projection dispatch " + "was not restored before the QKV resize hook ran"); + } + } + } + } +} + +} // namespace + +bool qkv_bk64_device_supported(WGPUDevice device) { + WGPULimits limits = {}; + const WebGPUContext* context = get_default_webgpu_context(); + return context != nullptr && context->shader_f16_supported && + wgpuDeviceGetLimits(device, &limits) == WGPUStatus_Success && + limits.maxComputeInvocationsPerWorkgroup >= 256u && + limits.maxComputeWorkgroupSizeX >= 16u && + limits.maxComputeWorkgroupSizeY >= 16u && + limits.maxComputeWorkgroupStorageSize >= 16384u && + limits.maxComputeWorkgroupsPerDimension >= 384u; +} + +void detect_qkv_bk64_fusions( + const WebGPUGraph& graph, + const vkgraph::VkGraph* fb_graph, + int num_vals, + std::vector& fusions, + std::unordered_map& first_ops, + std::unordered_map& last_ops, + std::unordered_map& member_ops) { + const auto* chain = fb_graph->chain(); + if (!chain || !qkv_bk64_device_supported(graph.device())) { + return; + } + std::unordered_map> q4_ops_by_input; + std::vector input_order; + for (unsigned i = 0; i < chain->size(); i++) { + const auto* op = chain->Get(i); + const auto* args = op->args(); + if (op->name()->str() != kQ4gswLinearOpName || !args || args->size() != 6) { + continue; + } + const int input_id = static_cast(args->Get(0)); + if (q4_ops_by_input.count(input_id) == 0) { + input_order.push_back(input_id); + } + q4_ops_by_input[input_id].push_back(i); + } + + auto op_arg = [&](unsigned op_index, unsigned arg_index) { + return static_cast(chain->Get(op_index)->args()->Get(arg_index)); + }; + const auto& output_ids = graph.output_ids(); + auto is_graph_output = [&](int id) { + return std::find(output_ids.begin(), output_ids.end(), id) != + output_ids.end(); + }; + for (int input_id : input_order) { + const auto& ops = q4_ops_by_input.at(input_id); + if (ops.size() != 3 || input_id < 0 || input_id >= num_vals || + graph.get_value_type(input_id) != WebGPUGraph::ValueType::Tensor) { + continue; + } + + QkvBk64Fusion fusion; + fusion.input_id = input_id; + bool exact_args = true; + for (size_t member = 0; member < 3; member++) { + fusion.op_indices[member] = ops[member]; + fusion.weight_ids[member] = op_arg(ops[member], 1); + fusion.scale_ids[member] = op_arg(ops[member], 2); + fusion.output_ids[member] = op_arg(ops[member], 5); + const int group_size_id = op_arg(ops[member], 3); + const int bias_id = op_arg(ops[member], 4); + exact_args = exact_args && group_size_id >= 0 && + group_size_id < num_vals && + graph.get_value_type(group_size_id) == WebGPUGraph::ValueType::Int && + graph.get_int(group_size_id) == kQkvGroupSize && bias_id >= 0 && + bias_id < num_vals && + graph.get_value_type(bias_id) == WebGPUGraph::ValueType::Null; + } + if (!exact_args) { + continue; + } + + const std::array constant_ids = { + fusion.weight_ids[0], + fusion.weight_ids[1], + fusion.weight_ids[2], + fusion.scale_ids[0], + fusion.scale_ids[1], + fusion.scale_ids[2]}; + const std::unordered_set distinct_constants( + constant_ids.begin(), constant_ids.end()); + bool direct_constants = distinct_constants.size() == constant_ids.size(); + for (int id : constant_ids) { + direct_constants = direct_constants && id >= 0 && id < num_vals && + graph.get_value_type(id) == WebGPUGraph::ValueType::Tensor && + graph.has_constant_source(id) && + graph.get_tensor(id).buffer != nullptr; + } + if (!direct_constants) { + continue; + } + + const std::unordered_set distinct_outputs = { + fusion.output_ids[0], fusion.output_ids[1], fusion.output_ids[2]}; + bool outputs_ok = distinct_outputs.size() == 3; + for (int id : fusion.output_ids) { + outputs_ok = outputs_ok && id >= 0 && id < num_vals && + graph.get_value_type(id) == WebGPUGraph::ValueType::Tensor && + graph.mem_obj_id(id) >= 0 && !is_graph_output(id) && + utils::is_fp32_tensor(graph.get_tensor(id)); + } + if (!outputs_ok) { + continue; + } + + const auto& input = graph.get_tensor(input_id); + if (!utils::is_fp32_tensor(input) || input.dims.empty() || + input.dims.back() != kQkvK) { + continue; + } + const uint64_t input_numel = utils::numel_of(input.dims); + if (input_numel % kQkvK != 0u || input_numel / kQkvK < 128u || + input_numel / kQkvK > UINT32_MAX) { + continue; + } + fusion.max_m = static_cast(input_numel / kQkvK); + + const uint32_t widths[3] = {kQkvQWidth, kQkvKvWidth, kQkvKvWidth}; + bool exact_geometry = true; + for (size_t member = 0; member < 3; member++) { + const auto& weight = graph.get_tensor(fusion.weight_ids[member]); + const auto& scale = graph.get_tensor(fusion.scale_ids[member]); + const auto& output = graph.get_tensor(fusion.output_ids[member]); + exact_geometry = + exact_geometry && weight.dims.size() == 2 && + weight.dims[0] == widths[member] && weight.dims[1] == kQkvKPacked && + weight.nbytes == static_cast(widths[member]) * kQkvKPacked && + scale.dims.size() == 2 && scale.dims[0] == kQkvNumGroups && + scale.dims[1] == widths[member] && utils::is_fp32_tensor(scale) && + output.dims.size() == input.dims.size() && + std::equal( + input.dims.begin(), input.dims.end() - 1, output.dims.begin()) && + output.dims.back() == widths[member] && + utils::numel_of(output.dims) == + static_cast(fusion.max_m) * widths[member]; + } + if (!exact_geometry) { + continue; + } + + const size_t fusion_index = fusions.size(); + fusions.push_back(fusion); + first_ops[ops[0]] = fusion_index; + last_ops[ops[2]] = fusion_index; + for (unsigned op : ops) { + member_ops[op] = fusion_index; + } + } +} + +void retain_unclaimed_qkv_fusions( + std::vector& fusions, + std::unordered_map& first_ops, + std::unordered_map& last_ops, + std::unordered_map& member_ops, + std::unordered_set& claimed_ops) { + std::vector retained_fusions; + first_ops.clear(); + last_ops.clear(); + member_ops.clear(); + for (QkvBk64Fusion& fusion : fusions) { + bool overlaps = false; + for (unsigned op : fusion.op_indices) { + overlaps = overlaps || claimed_ops.count(op) != 0; + } + if (overlaps) { + continue; + } + const size_t fusion_index = retained_fusions.size(); + retained_fusions.push_back(std::move(fusion)); + const QkvBk64Fusion& retained = retained_fusions.back(); + first_ops[retained.op_indices[0]] = fusion_index; + last_ops[retained.op_indices[2]] = fusion_index; + for (unsigned op : retained.op_indices) { + member_ops[op] = fusion_index; + claimed_ops.insert(op); + } + } + fusions = std::move(retained_fusions); +} + +void add_qkv_bk64_dispatch(WebGPUGraph& graph, QkvBk64Fusion& fusion) { + const auto& input = graph.get_tensor(fusion.input_id); + const auto& output_q = graph.get_tensor(fusion.output_ids[0]); + const auto& output_k = graph.get_tensor(fusion.output_ids[1]); + const auto& output_v = graph.get_tensor(fusion.output_ids[2]); + const auto& weight_q = graph.get_tensor(fusion.weight_ids[0]); + const auto& weight_k = graph.get_tensor(fusion.weight_ids[1]); + const auto& weight_v = graph.get_tensor(fusion.weight_ids[2]); + const auto& scale_q = graph.get_tensor(fusion.scale_ids[0]); + const auto& scale_k = graph.get_tensor(fusion.scale_ids[1]); + const auto& scale_v = graph.get_tensor(fusion.scale_ids[2]); + + const size_t weight_row_bytes = kQkvKPacked; + WGPUBuffer fused_weight = graph.create_scratch_buffer( + static_cast(kQkvFusedWidth) * weight_row_bytes); + WGPUBuffer fused_scales = graph.create_scratch_buffer( + static_cast(kQkvNumGroups) * kQkvFusedWidth * sizeof(float)); + + WGPUCommandEncoder encoder = + wgpuDeviceCreateCommandEncoder(graph.device(), nullptr); + wgpuCommandEncoderCopyBufferToBuffer( + encoder, + weight_q.buffer, + 0, + fused_weight, + 0, + static_cast(kQkvQWidth) * weight_row_bytes); + wgpuCommandEncoderCopyBufferToBuffer( + encoder, + weight_k.buffer, + 0, + fused_weight, + static_cast(kQkvQWidth) * weight_row_bytes, + static_cast(kQkvKvWidth) * weight_row_bytes); + wgpuCommandEncoderCopyBufferToBuffer( + encoder, + weight_v.buffer, + 0, + fused_weight, + static_cast(kQkvQWidth + kQkvKvWidth) * weight_row_bytes, + static_cast(kQkvKvWidth) * weight_row_bytes); + for (uint32_t group = 0; group < kQkvNumGroups; group++) { + const uint64_t destination = + static_cast(group) * kQkvFusedWidth * sizeof(float); + wgpuCommandEncoderCopyBufferToBuffer( + encoder, + scale_q.buffer, + static_cast(group) * kQkvQWidth * sizeof(float), + fused_scales, + destination, + static_cast(kQkvQWidth) * sizeof(float)); + wgpuCommandEncoderCopyBufferToBuffer( + encoder, + scale_k.buffer, + static_cast(group) * kQkvKvWidth * sizeof(float), + fused_scales, + destination + static_cast(kQkvQWidth) * sizeof(float), + static_cast(kQkvKvWidth) * sizeof(float)); + wgpuCommandEncoderCopyBufferToBuffer( + encoder, + scale_v.buffer, + static_cast(group) * kQkvKvWidth * sizeof(float), + fused_scales, + destination + + static_cast(kQkvQWidth + kQkvKvWidth) * sizeof(float), + static_cast(kQkvKvWidth) * sizeof(float)); + } + WGPUCommandBuffer command = wgpuCommandEncoderFinish(encoder, nullptr); + wgpuQueueSubmit(graph.queue(), 1, &command); + wgpuCommandBufferRelease(command); + wgpuCommandEncoderRelease(encoder); + + const QkvBk64Params params = { + fusion.max_m, + kQkvFusedWidth, + kQkvK, + kQkvKPacked, + kQkvGroupSize, + kQkvFusedWidth, + 0u, + 0u}; + WGPUBuffer params_buffer = graph.create_params_buffer(params); + WGPUBuffer bias_dummy = graph.create_scratch_buffer(4); + + const bool initially_active = is_qkv_bk64_live_m(fusion.max_m); + const uint32_t workgroups = + ((fusion.max_m + kQkvTile - 1u) / kQkvTile) * (kQkvFusedWidth / kQkvTile); + WebGPUComputeDispatchDescriptor descriptor; + descriptor.shader_name = "q4gsw_qkv_bk64"; + descriptor.kernel_name = "linear_q4gsw_bk64_qkv"; + descriptor.bindings = { + {output_q.buffer, 0u, output_q.nbytes}, + {output_k.buffer, 0u, output_k.nbytes}, + {output_v.buffer, 0u, output_v.nbytes}, + {input.buffer, 0u, input.nbytes}, + {fused_weight, + 0u, + static_cast(kQkvFusedWidth) * weight_row_bytes}, + {fused_scales, + 0u, + static_cast(kQkvNumGroups) * kQkvFusedWidth * sizeof(float)}, + {bias_dummy, 0u, 4u}, + {params_buffer, 0u, sizeof(QkvBk64Params)}}; + descriptor.grid = { + initially_active ? workgroups : 0u, initially_active ? 1u : 0u}; + fusion.fused_dispatch = graph.add_compute_dispatch(descriptor); + fusion.params_buffer = params_buffer; +} + +void add_qkv_bk64_resize_hook(WebGPUGraph& graph, const QkvBk64Fusion& fusion) { + const QkvBk64ResizeContext context = { + fusion.input_id, + {fusion.output_ids[0], fusion.output_ids[1], fusion.output_ids[2]}, + {fusion.separate_begin[0], + fusion.separate_begin[1], + fusion.separate_begin[2]}, + {fusion.separate_end[0], fusion.separate_end[1], fusion.separate_end[2]}, + fusion.fused_dispatch, + fusion.max_m, + fusion.params_buffer}; + resize_qkv_bk64(graph, context); + graph.add_tensor_resize_hook(fusion.input_id, [context](WebGPUGraph& g) { + resize_qkv_bk64(g, context); + }); +} + +} // namespace executorch::backends::webgpu::passes diff --git a/backends/webgpu/runtime/passes/QkvBk64.h b/backends/webgpu/runtime/passes/QkvBk64.h new file mode 100644 index 00000000000..5effc2ec04f --- /dev/null +++ b/backends/webgpu/runtime/passes/QkvBk64.h @@ -0,0 +1,71 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include +#include + +#include +#include +#include + +namespace executorch::backends::webgpu::passes { + +// One matched QKV-BK64 (three q4gsw linears sharing one input, exact Llama +// Q/K/V geometry) pattern. +struct QkvBk64Fusion { + int input_id = -1; + int output_ids[3] = {-1, -1, -1}; + int weight_ids[3] = {-1, -1, -1}; + int scale_ids[3] = {-1, -1, -1}; + unsigned op_indices[3] = {0, 0, 0}; + size_t separate_begin[3] = {0, 0, 0}; + size_t separate_end[3] = {0, 0, 0}; + size_t fused_dispatch = SIZE_MAX; + WGPUBuffer params_buffer = nullptr; + uint32_t max_m = 0; +}; + +// True if the device meets the BK64 kernel's shader-f16/workgroup limits. +bool qkv_bk64_device_supported(WGPUDevice device); + +// Phase 2: scan fb_graph's op chain for three q4gsw-linear ops sharing one +// input in exact Q/K/V geometry. Populates `fusions` and the per-op index +// maps Phase 3 uses; does NOT filter against already-claimed op indices -- +// SwiGLU keeps precedence over an overlapping QKV candidate, so call +// retain_unclaimed_qkv_fusions after SwiGLU detection completes. +void detect_qkv_bk64_fusions( + const WebGPUGraph& graph, + const vkgraph::VkGraph* fb_graph, + int num_vals, + std::vector& fusions, + std::unordered_map& first_ops, + std::unordered_map& last_ops, + std::unordered_map& member_ops); + +// Drops any QKV candidate overlapping an op index already in `claimed_ops` +// (claimed by a higher-precedence pass), rebuilds the index maps for the +// retained set, and adds the retained candidates' op indices to +// `claimed_ops`. +void retain_unclaimed_qkv_fusions( + std::vector& fusions, + std::unordered_map& first_ops, + std::unordered_map& last_ops, + std::unordered_map& member_ops, + std::unordered_set& claimed_ops); + +// Emits the single fused q4gsw_qkv_bk64 dispatch for a matched pattern. +void add_qkv_bk64_dispatch(WebGPUGraph& graph, QkvBk64Fusion& fusion); + +// Registers the dynamic-resize hook that switches the fusion between its +// fused and separate-projection dispatches as the live M crosses the BK64 +// kernel's supported shapes. +void add_qkv_bk64_resize_hook(WebGPUGraph& graph, const QkvBk64Fusion& fusion); + +} // namespace executorch::backends::webgpu::passes diff --git a/backends/webgpu/runtime/passes/SwiGLU.cpp b/backends/webgpu/runtime/passes/SwiGLU.cpp new file mode 100644 index 00000000000..4a98cfdb4f7 --- /dev/null +++ b/backends/webgpu/runtime/passes/SwiGLU.cpp @@ -0,0 +1,338 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include + +#include +#include + +#include +#include +#include +#include + +namespace executorch::backends::webgpu::passes { + +namespace { + +constexpr const char* kQ4gswLinearOpName = "et_vk.linear_q4gsw.default"; +constexpr const char* kSigmoidOpName = "aten.sigmoid.default"; +constexpr const char* kMulOpName = "aten.mul.Tensor"; + +// Uniform layout matching silu_mul_fused.wgsl's Params struct. +struct SiluMulParams { + uint32_t num_elements; + uint32_t _pad[3]; +}; + +uint32_t checked_silu_mul_numel(const std::vector& dims) { + const uint64_t numel = utils::numel_of(dims); + if (numel == 0 || numel > std::numeric_limits::max()) { + throw std::runtime_error("silu_mul_fused: element count out of range"); + } + return static_cast(numel); +} + +void resize_silu_mul_fused( + WebGPUGraph& graph, + int gate_id, + int up_id, + int out_id, + uint32_t wg_size, + size_t dispatch_idx, + WGPUBuffer params_buffer) { + const auto& gate_dims = graph.cur_dims(gate_id); + const auto& up_dims = graph.cur_dims(up_id); + if (gate_dims != up_dims) { + throw std::runtime_error("silu_mul_fused(resize): gate/up shape mismatch"); + } + const uint32_t live_numel = checked_silu_mul_numel(gate_dims); + const size_t live_nbytes = static_cast(live_numel) * sizeof(float); + if (graph.get_tensor(gate_id).cur_nbytes != live_nbytes || + graph.get_tensor(up_id).cur_nbytes != live_nbytes) { + throw std::runtime_error( + "silu_mul_fused(resize): gate/up byte-size mismatch"); + } + graph.set_cur_dims(out_id, gate_dims); + const SiluMulParams params = {live_numel, {0u, 0u, 0u}}; + wgpuQueueWriteBuffer( + graph.queue(), params_buffer, 0, ¶ms, sizeof(params)); + const utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( + graph.device(), live_numel, wg_size, "silu_mul_fused(resize)"); + auto& dispatch = graph.dispatch_at(dispatch_idx); + dispatch.workgroup_count_x = workgroup_count.x; + dispatch.workgroup_count_y = workgroup_count.y; +} + +} // namespace + +void detect_swiglu_fusions( + const WebGPUGraph& graph, + const vkgraph::VkGraph* fb_graph, + int num_vals, + std::vector& fusions, + std::unordered_map& gate_producers, + std::unordered_map& anchors, + std::unordered_set& skipped_ops, + std::unordered_set& claimed_ops) { + const auto* chain = fb_graph->chain(); + // Detect only the exact q4 gate/up MLP pattern. Full-chain occurrence and + // definition counts make every folded intermediate private to the pattern. + if (!chain) { + return; + } + struct ExactUnary { + unsigned op; + int input; + }; + struct ExactBinary { + unsigned op; + int lhs; + int rhs; + }; + std::vector occurrences(num_vals, 0); + std::vector definitions(num_vals, 0); + std::vector producer(num_vals, -1); + std::unordered_map q4_by_output; + std::unordered_map sigmoid_by_output; + std::unordered_map mul_by_output; + + auto count_occurrence = [&](int id) { + if (id < 0 || id >= num_vals) { + return; + } + occurrences[id]++; + if (graph.get_value_type(id) == WebGPUGraph::ValueType::ValueList) { + for (int member : graph.get_value_list(id)) { + if (member >= 0 && member < num_vals) { + occurrences[member]++; + } + } + } + }; + + for (unsigned i = 0; i < chain->size(); i++) { + const auto* op = chain->Get(i); + const auto* args = op->args(); + if (!args || args->size() == 0) { + continue; + } + for (unsigned j = 0; j < args->size(); j++) { + count_occurrence(static_cast(args->Get(j))); + } + const std::string name = op->name()->str(); + int output = -1; + if (name == kQ4gswLinearOpName && args->size() == 6) { + output = static_cast(args->Get(5)); + q4_by_output[output] = i; + } else if (name == kSigmoidOpName && args->size() == 2) { + output = static_cast(args->Get(1)); + sigmoid_by_output[output] = {i, static_cast(args->Get(0))}; + } else if (name == kMulOpName && args->size() == 3) { + output = static_cast(args->Get(2)); + mul_by_output[output] = { + i, static_cast(args->Get(0)), static_cast(args->Get(1))}; + } + if (output >= 0 && output < num_vals) { + definitions[output]++; + producer[output] = static_cast(i); + } + } + + const auto& output_ids = graph.output_ids(); + auto is_graph_output = [&](int id) { + return std::find(output_ids.begin(), output_ids.end(), id) != + output_ids.end(); + }; + for (unsigned mul2_op = 0; mul2_op < chain->size(); mul2_op++) { + const auto* mul2_call = chain->Get(mul2_op); + const auto* mul2_args = mul2_call->args(); + if (mul2_call->name()->str() != kMulOpName || !mul2_args || + mul2_args->size() != 3) { + continue; + } + + std::vector candidates; + auto try_orientation = [&](int silu_id, int up_id) { + const auto mul1_it = mul_by_output.find(silu_id); + if (mul1_it == mul_by_output.end()) { + return; + } + const ExactBinary& mul1 = mul1_it->second; + int gate_id = -1; + int sigmoid_id = -1; + unsigned sigmoid_op = 0; + const auto lhs_sig = sigmoid_by_output.find(mul1.lhs); + const auto rhs_sig = sigmoid_by_output.find(mul1.rhs); + if (lhs_sig != sigmoid_by_output.end() && + lhs_sig->second.input == mul1.rhs) { + gate_id = mul1.rhs; + sigmoid_id = mul1.lhs; + sigmoid_op = lhs_sig->second.op; + } else if ( + rhs_sig != sigmoid_by_output.end() && + rhs_sig->second.input == mul1.lhs) { + gate_id = mul1.lhs; + sigmoid_id = mul1.rhs; + sigmoid_op = rhs_sig->second.op; + } else { + return; + } + + const auto gate_q4 = q4_by_output.find(gate_id); + const auto up_q4 = q4_by_output.find(up_id); + if (gate_q4 == q4_by_output.end() || up_q4 == q4_by_output.end() || + gate_q4->second == up_q4->second) { + return; + } + const auto* gate_args = chain->Get(gate_q4->second)->args(); + const auto* up_args = chain->Get(up_q4->second)->args(); + if (!gate_args || !up_args || gate_args->size() != 6 || + up_args->size() != 6 || gate_args->Get(0) != up_args->Get(0) || + static_cast(gate_args->Get(5)) != gate_id || + static_cast(up_args->Get(5)) != up_id) { + return; + } + const int common_input_id = static_cast(gate_args->Get(0)); + const int out_id = static_cast(mul2_args->Get(2)); + const int ids[] = {gate_id, up_id, sigmoid_id, silu_id, out_id}; + std::unordered_set distinct_ids(std::begin(ids), std::end(ids)); + if (distinct_ids.size() != 5 || common_input_id < 0 || + common_input_id >= num_vals) { + return; + } + for (int id : ids) { + if (id < 0 || id >= num_vals || + graph.get_value_type(id) != WebGPUGraph::ValueType::Tensor || + definitions[id] != 1) { + return; + } + } + if (producer[gate_id] != static_cast(gate_q4->second) || + producer[up_id] != static_cast(up_q4->second) || + producer[sigmoid_id] != static_cast(sigmoid_op) || + producer[silu_id] != static_cast(mul1.op) || + producer[out_id] != static_cast(mul2_op) || + occurrences[gate_id] != 3 || occurrences[up_id] != 2 || + occurrences[sigmoid_id] != 2 || occurrences[silu_id] != 2) { + return; + } + if (!(gate_q4->second < sigmoid_op && sigmoid_op < mul1.op && + mul1.op < mul2_op && up_q4->second < mul2_op)) { + return; + } + if (is_graph_output(gate_id) || is_graph_output(sigmoid_id) || + is_graph_output(silu_id) || graph.mem_obj_id(gate_id) < 0) { + return; + } + + const auto& gate = graph.get_tensor(gate_id); + const auto& up = graph.get_tensor(up_id); + const auto& sigmoid = graph.get_tensor(sigmoid_id); + const auto& silu = graph.get_tensor(silu_id); + const auto& out = graph.get_tensor(out_id); + if (!utils::is_fp32_tensor(gate) || !utils::is_fp32_tensor(up) || + !utils::is_fp32_tensor(sigmoid) || !utils::is_fp32_tensor(silu) || + !utils::is_fp32_tensor(out) || gate.dims != up.dims || + gate.dims != sigmoid.dims || gate.dims != silu.dims || + gate.dims != out.dims || gate.nbytes != up.nbytes || + gate.nbytes != sigmoid.nbytes || gate.nbytes != silu.nbytes || + gate.nbytes != out.nbytes || up.buffer == out.buffer || + gate.buffer == out.buffer) { + return; + } + candidates.push_back( + {common_input_id, + gate_id, + up_id, + sigmoid_id, + silu_id, + out_id, + gate_q4->second, + sigmoid_op, + mul1.op, + mul2_op}); + }; + + try_orientation( + static_cast(mul2_args->Get(0)), + static_cast(mul2_args->Get(1))); + try_orientation( + static_cast(mul2_args->Get(1)), + static_cast(mul2_args->Get(0))); + if (candidates.size() != 1) { + continue; + } + const SwiGluFusion& fusion = candidates.front(); + const unsigned pattern_ops[] = { + fusion.gate_op, + q4_by_output.at(fusion.up_id), + fusion.sigmoid_op, + fusion.mul1_op, + fusion.mul2_op}; + bool overlaps = false; + for (unsigned op : pattern_ops) { + overlaps = overlaps || claimed_ops.count(op) != 0; + } + if (overlaps) { + continue; + } + const size_t fusion_idx = fusions.size(); + fusions.push_back(fusion); + gate_producers[fusion.gate_op] = fusion_idx; + anchors[fusion.mul2_op] = fusion_idx; + skipped_ops.insert(fusion.sigmoid_op); + skipped_ops.insert(fusion.mul1_op); + // mul2_op is the fusion anchor: its Phase-3 branch emits the fused + // dispatch and continues before the skipped-ops check, so it needs no + // skipped_ops entry. + for (unsigned op : pattern_ops) { + claimed_ops.insert(op); + } + } +} + +void add_silu_mul_fused_dispatch( + WebGPUGraph& graph, + int common_input_id, + int gate_id, + int up_id, + int out_id) { + const auto& gate = graph.get_tensor(gate_id); + const auto& up = graph.get_tensor(up_id); + const auto& out = graph.get_tensor(out_id); + const uint32_t num_elements = checked_silu_mul_numel(gate.dims); + const uint32_t wg_size = utils::clamp_workgroup_size( + graph.device(), + get_webgpu_shader_info("silu_mul_fused").workgroup_size_x); + const utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( + graph.device(), num_elements, wg_size, "silu_mul_fused"); + + const SiluMulParams params = {num_elements, {0u, 0u, 0u}}; + WGPUBuffer params_buffer = graph.create_params_buffer(params); + WebGPUComputeDispatchDescriptor descriptor; + descriptor.shader_name = "silu_mul_fused"; + descriptor.bindings = { + {gate.buffer, 0u, gate.nbytes}, + {up.buffer, 0u, up.nbytes}, + {out.buffer, 0u, out.nbytes}, + {params_buffer, 0u, sizeof(SiluMulParams)}}; + descriptor.constants = {{"wg_size", static_cast(wg_size)}}; + descriptor.grid = {workgroup_count.x, workgroup_count.y}; + const size_t dispatch_idx = graph.add_compute_dispatch(descriptor); + + graph.add_tensor_resize_hook( + common_input_id, + [gate_id, up_id, out_id, wg_size, dispatch_idx, params_buffer]( + WebGPUGraph& g) { + resize_silu_mul_fused( + g, gate_id, up_id, out_id, wg_size, dispatch_idx, params_buffer); + }); +} + +} // namespace executorch::backends::webgpu::passes diff --git a/backends/webgpu/runtime/passes/SwiGLU.h b/backends/webgpu/runtime/passes/SwiGLU.h new file mode 100644 index 00000000000..da0e2ca1b3d --- /dev/null +++ b/backends/webgpu/runtime/passes/SwiGLU.h @@ -0,0 +1,59 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include +#include + +#include +#include +#include + +namespace executorch::backends::webgpu::passes { + +// One matched SwiGLU (mul(mul(sigmoid(gate), gate), up) -> out) pattern. +struct SwiGluFusion { + int common_input_id; + int gate_id; + int up_id; + int sigmoid_id; + int silu_id; + int out_id; + unsigned gate_op; + unsigned sigmoid_op; + unsigned mul1_op; + unsigned mul2_op; +}; + +// Phase 2: scan fb_graph's op chain for the exact q4-gate/up + sigmoid + 2x +// mul SwiGLU pattern, skipping any op index already in `claimed_ops`. +// Populates `fusions` and the per-op index maps WebGPUGraph::build's Phase 3 +// walk uses to fire the fused dispatch, and adds every op index a matched +// pattern owns to `claimed_ops` so a later fusion pass (e.g. QKV-BK64) +// doesn't also try to claim them. +void detect_swiglu_fusions( + const WebGPUGraph& graph, + const vkgraph::VkGraph* fb_graph, + int num_vals, + std::vector& fusions, + std::unordered_map& gate_producers, + std::unordered_map& anchors, + std::unordered_set& skipped_ops, + std::unordered_set& claimed_ops); + +// Emits the single fused silu_mul_fused dispatch for a matched pattern and +// registers its dynamic-resize hook. +void add_silu_mul_fused_dispatch( + WebGPUGraph& graph, + int common_input_id, + int gate_id, + int up_id, + int out_id); + +} // namespace executorch::backends::webgpu::passes diff --git a/backends/webgpu/scripts/gen_wgsl_headers.py b/backends/webgpu/scripts/gen_wgsl_headers.py index c66aff2039a..1e35b205888 100644 --- a/backends/webgpu/scripts/gen_wgsl_headers.py +++ b/backends/webgpu/scripts/gen_wgsl_headers.py @@ -26,11 +26,14 @@ import copy import hashlib import io +import os import re +import stat import sys +import tempfile from itertools import product from pathlib import Path -from typing import Any, Dict, List, Optional, Set +from typing import Any, Dict, List, NamedTuple, Optional, Set, Tuple import yaml from yaml.constructor import ConstructorError @@ -124,6 +127,12 @@ def escape(line: str) -> str: def preprocess( input_text: str, variables: Dict[str, Any], input_path: str = "codegen" ) -> str: + # Normalize line endings first. Templates checked out with CRLF (common on + # Windows) otherwise break the trailing-backslash handling below: in + # re.MULTILINE, $ matches immediately before \n, so a CR would sit between a + # trailing \ and the line end and defeat the r"\\$" match, leaving a lone + # backslash that escape() turns into an unterminated Python string literal. + input_text = input_text.replace("\r\n", "\n").replace("\r", "\n") # Workaround to handle source files using \ to extend mecros to a new line input_text = re.sub(r"\\$", r"\\\\", input_text, flags=re.MULTILINE) @@ -431,14 +440,17 @@ def embedded_sha256(header_text: str) -> str: def _wg_size_const(base: str, axis: str, val: int) -> str: """One WorkgroupSize constant; wrap to <=80 cols so CLANGFORMAT accepts it. - Long shader names push the single-line form past the 80-col limit (clang-format - then breaks after '=' with a 4-space continuation indent); emit that wrapped - form up front so the generated header matches lintrunner's CLANGFORMAT. + Long shader names push the single-line form past the 80-col limit. Emit the + wrapped form that clang-format selects so generated headers stay byte-stable. """ - decl = f"inline constexpr uint32_t k{base}WorkgroupSize{axis} =" - if len(decl) + len(f" {val};") > 80: - return f"{decl}\n {val};\n" - return f"{decl} {val};\n" + name = f"k{base}WorkgroupSize{axis}" + prefix = f"inline constexpr uint32_t {name} =" + decl = f"{prefix} {val};" + if len(decl) > 85: + return f"inline constexpr uint32_t\n {name} = {val};\n" + if len(decl) > 80: + return f"{prefix}\n {val};\n" + return f"{decl}\n" def render_header( @@ -463,6 +475,14 @@ def render_header( raise ValueError('shader contains )" which would close the R"( literal') base = symbol_base(name) x, y, z = parse_workgroup_size(wgsl_text) + provenance = f"// @generated from {provenance_stem}.wgsl - DO NOT EDIT." + if len(provenance) > 80: + provenance_lines = [ + f"// @generated from {provenance_stem}.wgsl", + "// DO NOT EDIT.", + ] + else: + provenance_lines = [provenance] head = [ _BSD_HEADER, @@ -473,7 +493,7 @@ def render_header( "", "namespace executorch::backends::webgpu {", "", - f"// @generated from {provenance_stem}.wgsl - DO NOT EDIT.", + *provenance_lines, f"// wgsl-sha256: {wgsl_sha256(wgsl_text)}", f'inline constexpr const char* k{base}WGSL = R"(', ] @@ -496,6 +516,121 @@ def discover(): return sorted((BACKEND_ROOT / "runtime/ops").glob("**/*.wgsl")) +class RegistryEntry(NamedTuple): + name: str + include: str + symbol: str + + +def registry_path() -> Path: + return BACKEND_ROOT / "runtime/WebGPUShaderRegistry.cpp" + + +def _registry_entry(header: Path) -> RegistryEntry: + suffix = "_wgsl.h" + if not header.name.endswith(suffix): + raise ValueError(f"unexpected generated header name: {header.name}") + name = header.name[: -len(suffix)] + return RegistryEntry( + name=name, + include=header.relative_to(BACKEND_ROOT).as_posix(), + symbol=symbol_base(name), + ) + + +def _collect_header_outputs() -> Tuple[Dict[Path, str], List[RegistryEntry]]: + """Render every concrete header once and reject global collisions.""" + outputs: Dict[Path, str] = {} + entries: List[RegistryEntry] = [] + registry_names: Set[str] = set() + registry_symbols: Set[str] = set() + for wgsl in discover(): + try: + rendered_headers = list(headers_for_shader(wgsl)) + except Exception as error: + raise ValueError(f"{wgsl.relative_to(BACKEND_ROOT)}: {error}") from error + for header, rendered in rendered_headers: + if header in outputs: + raise ValueError( + "duplicate generated header path: " + f"{header.relative_to(BACKEND_ROOT)}" + ) + entry = _registry_entry(header) + if entry.name in registry_names: + raise ValueError(f"duplicate shader registry name: {entry.name}") + if entry.symbol in registry_symbols: + raise ValueError(f"duplicate shader registry symbol: {entry.symbol}") + outputs[header] = rendered + entries.append(entry) + registry_names.add(entry.name) + registry_symbols.add(entry.symbol) + return outputs, sorted(entries) + + +def registry_entries() -> List[RegistryEntry]: + """Return one registry entry for every concrete generated shader.""" + _, entries = _collect_header_outputs() + return entries + + +def render_registry(entries: List[RegistryEntry]) -> str: + """Render the generated name-to-WGSL registry implementation.""" + ordered = sorted(entries) + names = [entry.name for entry in ordered] + if len(names) != len(set(names)): + raise ValueError("duplicate shader registry name") + + includes = "\n".join( + sorted( + f"#include " + for entry in ordered + ) + ) + values = "\n".join( + " {\n" + f' "{entry.name}",\n' + f" k{entry.symbol}WGSL,\n" + f" k{entry.symbol}WorkgroupSizeX,\n" + f" k{entry.symbol}WorkgroupSizeY,\n" + f" k{entry.symbol}WorkgroupSizeZ,\n" + " }," + for entry in ordered + ) + return f"""{_BSD_HEADER} + +// @generated by scripts/gen_wgsl_headers.py - DO NOT EDIT. + +#include + +{includes} + +#include +#include +#include + +namespace executorch::backends::webgpu {{ +namespace {{ + +constexpr std::array kShaderRegistry = {{{{ +{values} +}}}}; + +}} // namespace + +const WebGPUShaderInfo& get_webgpu_shader_info(std::string_view name) {{ + for (const auto& shader : kShaderRegistry) {{ + if (shader.name == name) {{ + return shader; + }} + }} + throw std::runtime_error( + "WebGPU shader registry: unknown shader '" + std::string(name) + "'"); +}} + +}} // namespace executorch::backends::webgpu +""" + + def headers_for_shader(wgsl): """Yield (header_path, rendered_text) pairs for one shader source. @@ -526,7 +661,143 @@ def headers_for_shader(wgsl): yield header, render_header(stem, text, stem) -def _report_drift(missing, stale) -> None: +def collect_outputs() -> Tuple[Dict[Path, bytes], List[Path]]: + """Render the complete output tree and report unexpected old headers.""" + header_outputs, entries = _collect_header_outputs() + outputs = { + path: rendered.encode("utf-8") for path, rendered in header_outputs.items() + } + registry = registry_path() + if registry in outputs: + raise ValueError(f"duplicate generated output path: {registry}") + outputs[registry] = render_registry(entries).encode("utf-8") + + expected_headers = set(header_outputs) + actual_headers = set((BACKEND_ROOT / "runtime/ops").glob("**/*_wgsl.h")) + return outputs, sorted(actual_headers - expected_headers) + + +class _OriginalOutput(NamedTuple): + existed: bool + contents: bytes + mode: int + + +def _stage_bytes(destination: Path, contents: bytes, mode: int) -> Path: + """Write one same-directory candidate without changing its destination.""" + fd, name = tempfile.mkstemp( + prefix=f".{destination.name}.wgsl-gen-", + suffix=".tmp", + dir=destination.parent, + ) + temporary = Path(name) + try: + with os.fdopen(fd, "wb") as output: + output.write(contents) + temporary.chmod(mode) + except BaseException: + try: + temporary.unlink(missing_ok=True) + except OSError: + pass + raise + return temporary + + +def _cleanup_temporaries(temporaries) -> List[str]: + errors = [] + for temporary in temporaries: + try: + temporary.unlink(missing_ok=True) + except OSError as error: + errors.append(f"cannot remove temporary {temporary}: {error}") + return errors + + +def _stage_outputs( + outputs: Dict[Path, bytes], changed: List[Path] +) -> Tuple[Dict[Path, _OriginalOutput], Dict[Path, Path], List[str]]: + originals: Dict[Path, _OriginalOutput] = {} + staged: Dict[Path, Path] = {} + try: + for destination in sorted(changed): + if destination.exists(): + original = _OriginalOutput( + existed=True, + contents=destination.read_bytes(), + mode=stat.S_IMODE(destination.stat().st_mode), + ) + else: + original = _OriginalOutput(False, b"", 0o644) + originals[destination] = original + staged[destination] = _stage_bytes( + destination, outputs[destination], original.mode + ) + except BaseException as error: + cleanup_errors = _cleanup_temporaries(staged.values()) + if isinstance(error, OSError): + errors = [f"cannot stage generated output: {error}"] + cleanup_errors + return originals, staged, errors + raise + return originals, staged, [] + + +def _rollback_outputs( + originals: Dict[Path, _OriginalOutput], + replaced: List[Path], + staged: Dict[Path, Path], +) -> List[str]: + errors = [] + for destination in reversed(replaced): + original = originals[destination] + restore_temporary: Optional[Path] = None + try: + if original.existed: + restore_temporary = _stage_bytes( + destination, original.contents, original.mode + ) + os.replace(restore_temporary, destination) + else: + destination.unlink(missing_ok=True) + except OSError as error: + errors.append(f"cannot roll back {destination}: {error}") + finally: + if restore_temporary is not None: + errors.extend(_cleanup_temporaries([restore_temporary])) + errors.extend(_cleanup_temporaries(staged.values())) + return errors + + +def _publish_outputs(outputs: Dict[Path, bytes], changed: List[Path]) -> List[str]: + """Stage and publish changed outputs, rolling back reported failures.""" + originals, staged, stage_errors = _stage_outputs(outputs, changed) + if stage_errors: + return stage_errors + + replaced: List[Path] = [] + try: + for destination in sorted(changed): + try: + os.replace(staged[destination], destination) + except OSError: + raise + except BaseException: + replaced.append(destination) + raise + else: + replaced.append(destination) + except OSError as commit_error: + return [f"cannot publish generated output: {commit_error}"] + _rollback_outputs( + originals, replaced, staged + ) + except BaseException: + _rollback_outputs(originals, replaced, staged) + raise + + return _cleanup_temporaries(staged.values()) + + +def _report_drift(missing, stale, orphans) -> None: """Print the --check report for missing/stale committed headers.""" if missing: print("Missing embedded WGSL headers (run scripts/gen_wgsl_headers.py):") @@ -536,6 +807,10 @@ def _report_drift(missing, stale) -> None: print("Stale embedded WGSL headers (run scripts/gen_wgsl_headers.py):") for h in stale: print(f" {h.relative_to(BACKEND_ROOT)}") + if orphans: + print("Orphan embedded WGSL headers (remove or restore their sources):") + for h in orphans: + print(f" {h.relative_to(BACKEND_ROOT)}") def main(argv=None) -> int: @@ -547,35 +822,35 @@ def main(argv=None) -> int: ) args = parser.parse_args(argv) - stale = [] - missing = [] - errors = [] - for wgsl in discover(): - try: - rendered = list(headers_for_shader(wgsl)) - # A malformed spec raises yaml.YAMLError (incl. UniqueKeyLoader's - # ConstructorError) / ValueError / KeyError from parse_template_spec, and - # a malformed template raises AssertionError from preprocess; catch them - # all so a bad shader is a clean --check report, not a traceback. - except (ValueError, KeyError, AssertionError, yaml.YAMLError) as e: - errors.append(f"{wgsl.relative_to(BACKEND_ROOT)}: {e}") - continue - for header, want in rendered: - # Full-content compare (not just the sha) catches generator-logic drift too. - if header.exists() and header.read_text() == want: - continue - if args.check: - (missing if not header.exists() else stale).append(header) - else: - header.write_text(want) + try: + outputs, orphans = collect_outputs() + missing = [] + stale = [] + for output, want in sorted(outputs.items()): + if not output.exists(): + missing.append(output) + elif output.read_bytes() != want: + stale.append(output) + except Exception as error: + print("Cannot generate WGSL outputs:") + print(f" {error}") + return 1 - if errors: - print("Cannot generate header (malformed shader):") - for e in errors: - print(f" {e}") + if orphans: + _report_drift([], [], orphans) return 1 - if args.check and (stale or missing): - _report_drift(missing, stale) + + if args.check: + if stale or missing: + _report_drift(missing, stale, []) + return 1 + return 0 + + errors = _publish_outputs(outputs, missing + stale) + if errors: + print("Cannot publish WGSL outputs:") + for error in errors: + print(f" {error}") return 1 return 0 diff --git a/backends/webgpu/scripts/test_webgpu_native_ci.sh b/backends/webgpu/scripts/test_webgpu_native_ci.sh index 810d8165303..0d5d7a9e799 100644 --- a/backends/webgpu/scripts/test_webgpu_native_ci.sh +++ b/backends/webgpu/scripts/test_webgpu_native_ci.sh @@ -17,10 +17,8 @@ # source .ci/scripts/setup-webgpu-linux-deps.sh # bash backends/webgpu/scripts/test_webgpu_native_ci.sh # -# Builds whatever native test targets are present in the landed tree (NOT a fixed -# list): webgpu_native_test (base) + webgpu_dispatch_order_test, -# webgpu_scratch_buffer_test (D107576199) + webgpu_update_cache_test -# (D107547307). SDPA executables join once they land. +# Builds and runs the fixed native target matrix defined by this tree. A missing +# target or fixture is a CI failure, not an optional skip. set -e @@ -37,18 +35,39 @@ fi cd "${EXECUTORCH_ROOT}" -# ── Exports for the model-driven executables (best-effort) ─────────────────── -# native_test (quantized_linear/SDPA/update_cache) + dispatch_order read .pte/ -# golden inputs via env/dir and self-skip if absent; scratch is standalone. -# native_test itself is gated below on the executorch wheel being importable. +# ── Exports for the model-driven executables ───────────────────────────────── +if ! "${PYTHON_EXECUTABLE}" -c "import executorch" 2>/dev/null; then + echo "ERROR: executorch wheel unavailable; required fixture exports cannot run" >&2 + exit 1 +fi + +require_file() { + if [[ ! -f "$1" ]]; then + echo "ERROR: required WebGPU fixture missing: $1" >&2 + exit 1 + fi +} + +run_with_required_device() { + local output + if ! output="$("$@" 2>&1)"; then + printf '%s\n' "${output}" + return 1 + fi + printf '%s\n' "${output}" + if ! grep -q '^WebGPU device acquired (native)$' <<<"${output}"; then + echo "ERROR: WebGPU native test did not acquire a device" >&2 + return 1 + fi +} + DISPATCH_ORDER_DIR="/tmp/dispatch_order" -DISPATCH_ORDER_OK=1 UPDATE_CACHE_DIR="/tmp/update_cache" -UPDATE_CACHE_OK=1 INDEX_DIR="/tmp/index" -INDEX_OK=1 DYNAMIC_SHAPE_DIR="/tmp/dynamic_shape" -DYNAMIC_SHAPE_OK=1 +ROPE_HF_DIR="/tmp/webgpu_rope_hf" +SYMINT_BLOB="/tmp/sdpa_dyn_small.pte" +OUTPUT_SUPPRESSION_DIR="/tmp/output_suppression" EMBEDDING_MODEL="/tmp/webgpu_embedding_q4gsw.pte" EMBEDDING_INDICES="/tmp/webgpu_embedding_q4gsw_indices.bin" EMBEDDING_GOLDEN="/tmp/webgpu_embedding_q4gsw_golden.bin" @@ -69,33 +88,39 @@ PREPACK_TIED_MODEL="/tmp/webgpu_prepack_tied_const.pte" PREPACK_TIED_GOLDEN="/tmp/webgpu_prepack_tied_const_golden.bin" $PYTHON_EXECUTABLE -c " -from executorch.backends.webgpu.test.ops.test_quantized_linear import export_all_quantized_linear_models +from executorch.backends.webgpu.test.ops.test_quantized_linear import export_all_quantized_linear_models, export_output_suppression_models export_all_quantized_linear_models('/tmp') -" || echo "WARN: q4gsw export failed; required configs will FAIL in webgpu_native_test" +export_output_suppression_models('${OUTPUT_SUPPRESSION_DIR}') +" $PYTHON_EXECUTABLE -c " from executorch.backends.webgpu.test.ops.test_embedding_q4gsw import export_embedding_q4gsw_model export_embedding_q4gsw_model('${EMBEDDING_MODEL}', '${EMBEDDING_GOLDEN}', '${EMBEDDING_INDICES}') export_embedding_q4gsw_model('${EMBEDDING_LLAMA1B_MODEL}', '${EMBEDDING_LLAMA1B_GOLDEN}', '${EMBEDDING_LLAMA1B_INDICES}', 'llama1b') -" || echo "WARN: embedding_q4gsw export failed; embedding configs will FAIL in webgpu_native_test" +" $PYTHON_EXECUTABLE -c " from executorch.backends.webgpu.test.ops.test_rope import export_rope_model export_rope_model('${ROPE_MODEL}', '${ROPE_XQ_GOLDEN}', '${ROPE_XK_GOLDEN}') export_rope_model('${ROPE_DECODE_MODEL}', '${ROPE_DECODE_XQ_GOLDEN}', '${ROPE_DECODE_XK_GOLDEN}', 'decode') -" || echo "WARN: rope export failed; apply_rotary_emb configs will FAIL in webgpu_native_test" +" + +$PYTHON_EXECUTABLE -c " +from executorch.backends.webgpu.test.ops.test_rope_hf import export_rope_hf_dynamic +export_rope_hf_dynamic('${ROPE_HF_DIR}') +" $PYTHON_EXECUTABLE -c " from executorch.backends.webgpu.test.ops.test_prepack import export_prepack_model, export_prepack_two_const_model, export_prepack_tied_const_model export_prepack_model('${PREPACK_MODEL}', '${PREPACK_GOLDEN}') export_prepack_two_const_model('${PREPACK2_MODEL}', '${PREPACK2_GOLDEN}') export_prepack_tied_const_model('${PREPACK_TIED_MODEL}', '${PREPACK_TIED_GOLDEN}') -" || echo "WARN: prepack export failed; prepack configs will FAIL in webgpu_native_test" +" $PYTHON_EXECUTABLE -c " from executorch.backends.webgpu.test.ops.test_dispatch_order import export_dispatch_order_cases export_dispatch_order_cases('${DISPATCH_ORDER_DIR}') -" || { echo "WARN: dispatch_order export failed; skipping dispatch_order native test"; DISPATCH_ORDER_OK=0; } +" $PYTHON_EXECUTABLE -c " from executorch.backends.webgpu.test.ops.test_update_cache import ( @@ -106,20 +131,18 @@ from executorch.backends.webgpu.test.ops.test_update_cache import ( export_update_cache_cases('${UPDATE_CACHE_DIR}') export_update_cache_replay('${UPDATE_CACHE_DIR}') export_update_cache_negative('${UPDATE_CACHE_DIR}') -" || { echo "WARN: update_cache export failed; skipping update_cache native test"; UPDATE_CACHE_OK=0; } +" $PYTHON_EXECUTABLE -c " from executorch.backends.webgpu.test.ops.index.test_index import export_all_index_models export_all_index_models('${INDEX_DIR}') -" || { echo "WARN: index export failed; skipping index native test"; INDEX_OK=0; } +" $PYTHON_EXECUTABLE -c " from executorch.backends.webgpu.test.ops.dynamic_shape.test_dynamic_shape_export import export_dynamic_shape_cases export_dynamic_shape_cases('${DYNAMIC_SHAPE_DIR}') -" || { echo "WARN: dynamic_shape export failed; skipping dynamic_shape native test"; DYNAMIC_SHAPE_OK=0; } +" -# Non-fatal: a failed sdpa export makes the required 4k/8k configs hard-fail in -# webgpu_native_test below (precise per-config error), so don't exit/mask here. $PYTHON_EXECUTABLE -c " from executorch.backends.webgpu.test.ops.test_sdpa import ( export_all_sdpa_models, @@ -131,12 +154,17 @@ export_all_sdpa_models('/tmp') export_replay_sequences('/tmp') export_dynamic_decode('/tmp') export_incache_decode('/tmp') -" || echo "WARN: sdpa export failed; required 4k/8k configs will FAIL in webgpu_native_test" +" + +require_file "${ROPE_HF_DIR}/rope_hf_dynamic.pte" +require_file "${SYMINT_BLOB}" +require_file "${OUTPUT_SUPPRESSION_DIR}/input.bin" # ── Configure (Dawn-only: no -DWEBGPU_IMPL; Dawn is the sole backend) ───────── echo "=== Configure WebGPU native tests on Dawn ===" rm -rf "${BUILD_DIR}" cmake \ + -DPYTHON_EXECUTABLE="${PYTHON_EXECUTABLE}" \ -DEXECUTORCH_BUILD_WEBGPU=ON \ -DEXECUTORCH_BUILD_WEBGPU_TEST=ON \ -DEXECUTORCH_BUILD_TESTS=ON \ @@ -150,102 +178,64 @@ cmake \ -B "${BUILD_DIR}" \ "${EXECUTORCH_ROOT}" -# ── Build + run every native test target that exists in this tree ──────────── -TARGETS=(webgpu_native_test webgpu_dispatch_order_test webgpu_scratch_buffer_test webgpu_update_cache_test webgpu_index_test webgpu_dynamic_shape_test webgpu_dispatch_2d_test) +# ── Build + run every fixed native test target in this tree ────────────────── +REQUIRED_TARGETS=(webgpu_native_test webgpu_dispatch_order_test webgpu_scratch_buffer_test webgpu_update_cache_test webgpu_index_test webgpu_dynamic_shape_test webgpu_dispatch_2d_test webgpu_compute_dispatch_test webgpu_execution_options_test webgpu_output_suppression_test webgpu_op_test_util_test) BIN_DIR="${BUILD_DIR}/backends/webgpu" -# Which targets are defined depends on which diffs are landed (native_test + -# rms_norm here; dispatch_order + scratch from D107576199). Query the configured -# target list ONCE so a not-yet-landed target is skipped WITHOUT masking a real -# compile failure of a target that IS defined (CI uses the Make generator). DEFINED_TARGETS="$(cmake --build "${BUILD_DIR}" --target help 2>/dev/null || true)" -# Fail loud if the probe found nothing (e.g. a non-Make generator or a cmake -# regression): otherwise every target would skip and the job would go green -# having tested nothing. webgpu_native_test is always defined at/after this diff. -if ! printf '%s\n' "${DEFINED_TARGETS}" | grep -qw webgpu_native_test; then - echo "ERROR: cmake target probe returned no webgpu_native_test; aborting" >&2 - exit 1 -fi - -for t in "${TARGETS[@]}"; do - if printf '%s\n' "${DEFINED_TARGETS}" | grep -qw "${t}"; then - # Defined target: build with stderr visible; set -e fails the job on a real - # build error (never silently skipped). - cmake --build "${BUILD_DIR}" --target "${t}" -j"${NPROC}" - echo "built ${t}" - else - echo "(target ${t} not defined in this tree — skipping)" +for t in "${REQUIRED_TARGETS[@]}"; do + if ! printf '%s\n' "${DEFINED_TARGETS}" | grep -qw "${t}"; then + echo "ERROR: required CMake target is not defined: ${t}" >&2 + exit 1 fi + cmake --build "${BUILD_DIR}" --target "${t}" -j"${NPROC}" + echo "built ${t}" done echo "=== Run native tests on Dawn + SwiftShader ===" -# webgpu_native_test hosts the quantized_linear / SDPA / update_cache / symint -# sweeps. Gate on the executorch wheel being importable (the proxy for "the -# exports above ran"): CI has the wheel so they ran; a bare local run without it -# skips here rather than hard-failing the required-config guards. -if [[ -x "${BIN_DIR}/webgpu_native_test" ]] && - "${PYTHON_EXECUTABLE}" -c "import executorch" 2>/dev/null; then - env WEBGPU_TEST_SDPA_DIR=/tmp/ \ - WEBGPU_TEST_QUANTIZED_LINEAR_DIR=/tmp/ \ - WEBGPU_TEST_EMBEDDING_Q4GSW_MODEL="${EMBEDDING_MODEL}" \ - WEBGPU_TEST_EMBEDDING_Q4GSW_INDICES="${EMBEDDING_INDICES}" \ - WEBGPU_TEST_EMBEDDING_Q4GSW_GOLDEN="${EMBEDDING_GOLDEN}" \ - WEBGPU_TEST_EMBEDDING_Q4GSW_LLAMA1B_MODEL="${EMBEDDING_LLAMA1B_MODEL}" \ - WEBGPU_TEST_EMBEDDING_Q4GSW_LLAMA1B_INDICES="${EMBEDDING_LLAMA1B_INDICES}" \ - WEBGPU_TEST_EMBEDDING_Q4GSW_LLAMA1B_GOLDEN="${EMBEDDING_LLAMA1B_GOLDEN}" \ - WEBGPU_TEST_ROPE_MODEL="${ROPE_MODEL}" \ - WEBGPU_TEST_ROPE_XQ_GOLDEN="${ROPE_XQ_GOLDEN}" \ - WEBGPU_TEST_ROPE_XK_GOLDEN="${ROPE_XK_GOLDEN}" \ - WEBGPU_TEST_ROPE_DECODE_MODEL="${ROPE_DECODE_MODEL}" \ - WEBGPU_TEST_ROPE_DECODE_XQ_GOLDEN="${ROPE_DECODE_XQ_GOLDEN}" \ - WEBGPU_TEST_ROPE_DECODE_XK_GOLDEN="${ROPE_DECODE_XK_GOLDEN}" \ - WEBGPU_TEST_PREPACK_MODEL="${PREPACK_MODEL}" \ - WEBGPU_TEST_PREPACK_GOLDEN="${PREPACK_GOLDEN}" \ - WEBGPU_TEST_PREPACK2_MODEL="${PREPACK2_MODEL}" \ - WEBGPU_TEST_PREPACK2_GOLDEN="${PREPACK2_GOLDEN}" \ - WEBGPU_TEST_PREPACK_TIED_MODEL="${PREPACK_TIED_MODEL}" \ - WEBGPU_TEST_PREPACK_TIED_GOLDEN="${PREPACK_TIED_GOLDEN}" \ - "${BIN_DIR}/webgpu_native_test" -else - echo "(skipping webgpu_native_test: executorch wheel absent — exports did not run)" -fi -if [[ "${UPDATE_CACHE_OK}" == "1" && -x "${BIN_DIR}/webgpu_update_cache_test" ]]; then - "${BIN_DIR}/webgpu_update_cache_test" "${UPDATE_CACHE_DIR}" -fi -if [[ "${DISPATCH_ORDER_OK}" == "1" && -x "${BIN_DIR}/webgpu_dispatch_order_test" ]]; then - "${BIN_DIR}/webgpu_dispatch_order_test" "${DISPATCH_ORDER_DIR}" -fi -if [[ "${INDEX_OK}" == "1" && -x "${BIN_DIR}/webgpu_index_test" ]]; then - "${BIN_DIR}/webgpu_index_test" "${INDEX_DIR}" -fi -if [[ "${DYNAMIC_SHAPE_OK}" == "1" && -x "${BIN_DIR}/webgpu_dynamic_shape_test" ]]; then - "${BIN_DIR}/webgpu_dynamic_shape_test" "${DYNAMIC_SHAPE_DIR}" -fi -[[ -x "${BIN_DIR}/webgpu_scratch_buffer_test" ]] && "${BIN_DIR}/webgpu_scratch_buffer_test" -# Device-free: pure 2D workgroup-count fold unit test (no .pte, no GPU). -[[ -x "${BIN_DIR}/webgpu_dispatch_2d_test" ]] && "${BIN_DIR}/webgpu_dispatch_2d_test" +run_with_required_device env WEBGPU_TEST_SDPA_DIR=/tmp/ \ + WEBGPU_TEST_QUANTIZED_LINEAR_DIR=/tmp/ \ + WEBGPU_TEST_EMBEDDING_Q4GSW_MODEL="${EMBEDDING_MODEL}" \ + WEBGPU_TEST_EMBEDDING_Q4GSW_INDICES="${EMBEDDING_INDICES}" \ + WEBGPU_TEST_EMBEDDING_Q4GSW_GOLDEN="${EMBEDDING_GOLDEN}" \ + WEBGPU_TEST_EMBEDDING_Q4GSW_LLAMA1B_MODEL="${EMBEDDING_LLAMA1B_MODEL}" \ + WEBGPU_TEST_EMBEDDING_Q4GSW_LLAMA1B_INDICES="${EMBEDDING_LLAMA1B_INDICES}" \ + WEBGPU_TEST_EMBEDDING_Q4GSW_LLAMA1B_GOLDEN="${EMBEDDING_LLAMA1B_GOLDEN}" \ + WEBGPU_TEST_ROPE_MODEL="${ROPE_MODEL}" \ + WEBGPU_TEST_ROPE_XQ_GOLDEN="${ROPE_XQ_GOLDEN}" \ + WEBGPU_TEST_ROPE_XK_GOLDEN="${ROPE_XK_GOLDEN}" \ + WEBGPU_TEST_ROPE_DECODE_MODEL="${ROPE_DECODE_MODEL}" \ + WEBGPU_TEST_ROPE_DECODE_XQ_GOLDEN="${ROPE_DECODE_XQ_GOLDEN}" \ + WEBGPU_TEST_ROPE_DECODE_XK_GOLDEN="${ROPE_DECODE_XK_GOLDEN}" \ + WEBGPU_TEST_ROPE_HF_DIR="${ROPE_HF_DIR}" \ + WEBGPU_TEST_SYMINT_BLOB="${SYMINT_BLOB}" \ + WEBGPU_TEST_PREPACK_MODEL="${PREPACK_MODEL}" \ + WEBGPU_TEST_PREPACK_GOLDEN="${PREPACK_GOLDEN}" \ + WEBGPU_TEST_PREPACK2_MODEL="${PREPACK2_MODEL}" \ + WEBGPU_TEST_PREPACK2_GOLDEN="${PREPACK2_GOLDEN}" \ + WEBGPU_TEST_PREPACK_TIED_MODEL="${PREPACK_TIED_MODEL}" \ + WEBGPU_TEST_PREPACK_TIED_GOLDEN="${PREPACK_TIED_GOLDEN}" \ + "${BIN_DIR}/webgpu_native_test" +"${BIN_DIR}/webgpu_update_cache_test" "${UPDATE_CACHE_DIR}" +"${BIN_DIR}/webgpu_dispatch_order_test" "${DISPATCH_ORDER_DIR}" +"${BIN_DIR}/webgpu_index_test" "${INDEX_DIR}" +"${BIN_DIR}/webgpu_dynamic_shape_test" "${DYNAMIC_SHAPE_DIR}" +"${BIN_DIR}/webgpu_scratch_buffer_test" +"${BIN_DIR}/webgpu_dispatch_2d_test" +"${BIN_DIR}/webgpu_compute_dispatch_test" +"${BIN_DIR}/webgpu_execution_options_test" +"${BIN_DIR}/webgpu_output_suppression_test" "${OUTPUT_SUPPRESSION_DIR}" +"${BIN_DIR}/webgpu_op_test_util_test" echo "=== WebGPU native tests on Dawn: all run targets passed ===" # ── Op-test codegen framework: generate manifest → build → run (Dawn+SwiftShader) ── -# Reconfigure the SAME build dir adding GTest (EXECUTORCH_BUILD_TESTS=ON), then run -# every op in cases.py against its torch golden. Self-skips if the generator can't run. +# Generate the op-test manifest, build the target from the existing test-enabled +# configuration, then run every op in cases.py against its torch golden. OP_TEST_DIR="/tmp/webgpu_op_tests" -if $PYTHON_EXECUTABLE -m executorch.backends.webgpu.test.op_tests.generate_op_tests \ - --output "${OP_TEST_DIR}"; then - echo "=== Reconfigure with GTest + build/run op-test framework ===" - cmake -DEXECUTORCH_BUILD_TESTS=ON -B "${BUILD_DIR}" "${EXECUTORCH_ROOT}" - OP_DEFINED="$(cmake --build "${BUILD_DIR}" --target help 2>/dev/null || true)" - if printf '%s\n' "${OP_DEFINED}" | grep -qw webgpu_op_test_util_test; then - cmake --build "${BUILD_DIR}" --target webgpu_op_test_util_test -j"${NPROC}" - "${BIN_DIR}/webgpu_op_test_util_test" - fi - if printf '%s\n' "${OP_DEFINED}" | grep -qw webgpu_op_test; then - cmake --build "${BUILD_DIR}" --target webgpu_op_test -j"${NPROC}" - "${BIN_DIR}/webgpu_op_test" --manifest "${OP_TEST_DIR}/manifest.json" - fi - echo "=== WebGPU op-test framework on Dawn: passed ===" -else - echo "WARN: op-test manifest generation failed (needs the executorch wheel); skipping" -fi +$PYTHON_EXECUTABLE -m executorch.backends.webgpu.test.op_tests.generate_op_tests \ + --output "${OP_TEST_DIR}" +cmake --build "${BUILD_DIR}" --target webgpu_op_test -j"${NPROC}" +"${BIN_DIR}/webgpu_op_test" --manifest "${OP_TEST_DIR}/manifest.json" +echo "=== WebGPU op-test framework on Dawn: passed ===" diff --git a/backends/webgpu/test/native/test_compute_dispatch.cpp b/backends/webgpu/test/native/test_compute_dispatch.cpp new file mode 100644 index 00000000000..6ed7229604b --- /dev/null +++ b/backends/webgpu/test/native/test_compute_dispatch.cpp @@ -0,0 +1,947 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include + +namespace executorch::backends::webgpu { +namespace { + +// NOLINTBEGIN(cppcoreguidelines-avoid-non-const-global-variables) +WGPUDevice g_device = nullptr; +// NOLINTEND(cppcoreguidelines-avoid-non-const-global-variables) + +struct UnaryParams { + uint32_t num_elements; + float min; + float max; + uint32_t padding; +}; + +WGPUBuffer create_storage_buffer(size_t nbytes) { + WGPUBufferDescriptor descriptor = {}; + descriptor.size = nbytes; + descriptor.usage = WGPUBufferUsage_Storage | WGPUBufferUsage_CopyDst; + WGPUBuffer buffer = wgpuDeviceCreateBuffer(g_device, &descriptor); + if (buffer == nullptr) { + throw std::runtime_error("failed to create compute-dispatch test buffer"); + } + return buffer; +} + +enum class Q4RouteSignal { Static, LegacyGraphMarker, ExplicitOption }; + +void build_q4_route_graph(WebGPUGraph& graph, Q4RouteSignal signal) { + namespace vk = vkgraph; + ::flatbuffers::FlatBufferBuilder fbb; + std::vector<::flatbuffers::Offset> values; + auto add_tensor = [&](vk::VkDataType dtype, + const std::vector& dims, + int mem_obj_id) { + const int id = static_cast(values.size()); + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::VkTensor, + vk::CreateVkTensorDirect( + fbb, dtype, &dims, /*constant_id=*/-1, mem_obj_id) + .Union())); + return id; + }; + auto add_int = [&](int64_t value) { + const int id = static_cast(values.size()); + values.push_back(vk::CreateVkValue( + fbb, vk::GraphTypes::Int, vk::CreateInt(fbb, value).Union())); + return id; + }; + + const int input = add_tensor(vk::VkDataType::FLOAT32, {2, 8}, 0); + const int weight = add_tensor(vk::VkDataType::UINT8, {8, 4}, 1); + const int scales = add_tensor(vk::VkDataType::FLOAT32, {1, 8}, 2); + const int group_size = add_int(8); + const int bias = static_cast(values.size()); + values.push_back(vk::CreateVkValue(fbb)); + const int output = add_tensor(vk::VkDataType::FLOAT32, {2, 8}, 3); + + std::vector<::flatbuffers::Offset> chain; + int q4_input = input; + if (signal == Q4RouteSignal::LegacyGraphMarker) { + const int dim = add_int(0); + const int symint = static_cast(values.size()); + values.push_back(vk::CreateVkValue( + fbb, vk::GraphTypes::SymInt, vk::CreateSymInt(fbb, 0).Union())); + const int alpha = add_int(1); + const int intermediate = add_tensor(vk::VkDataType::FLOAT32, {2, 8}, 4); + const std::vector sym_size_args = {input, dim, symint}; + chain.push_back( + vk::CreateOperatorCallDirect(fbb, 0, "sym_size.int", &sym_size_args)); + const std::vector add_args = {input, input, alpha, intermediate}; + chain.push_back( + vk::CreateOperatorCallDirect(fbb, 1, "aten.add.Tensor", &add_args)); + q4_input = intermediate; + } + + const std::vector q4_args = { + q4_input, weight, scales, group_size, bias, output}; + chain.push_back(vk::CreateOperatorCallDirect( + fbb, + static_cast(chain.size()), + "et_vk.linear_q4gsw.default", + &q4_args)); + const std::vector input_ids = { + static_cast(input), + static_cast(weight), + static_cast(scales)}; + const std::vector output_ids = {static_cast(output)}; + const auto root = vk::CreateVkGraphDirect( + fbb, "0", &chain, &values, &input_ids, &output_ids); + vk::FinishVkGraphBuffer(fbb, root); + + WebGPUGraphConfig config; + if (signal == Q4RouteSignal::ExplicitOption) { + config.record_q4gsw_decode_route = true; + } + graph.build(fbb.GetBufferPointer(), nullptr, 0, nullptr, config); +} + +std::vector q4_dispatches(WebGPUGraph& graph) { + std::vector dispatches; + for (size_t i = 0; i < graph.num_dispatches(); i++) { + const WebGPUDispatch& dispatch = graph.dispatch_at(i); + if (dispatch.kernel_name.rfind("linear_q4gsw", 0) == 0) { + dispatches.push_back(&dispatch); + } + } + return dispatches; +} + +void expect_dual_q4_topology(WebGPUGraph& graph) { + const auto dispatches = q4_dispatches(graph); + ASSERT_EQ(dispatches.size(), 2); + EXPECT_NE(dispatches[0]->pipeline, nullptr); + EXPECT_NE(dispatches[1]->pipeline, nullptr); + EXPECT_NE(dispatches[0]->pipeline, dispatches[1]->pipeline); + EXPECT_NE(dispatches[0]->bind_group, nullptr); + EXPECT_EQ(dispatches[0]->bind_group, dispatches[1]->bind_group); + EXPECT_EQ( + std::count_if( + dispatches.begin(), + dispatches.end(), + [](const WebGPUDispatch* dispatch) { + return dispatch->workgroup_count_x != 0 && + dispatch->workgroup_count_y != 0; + }), + 1); +} + +TEST(WebGPUShaderRegistry, FindsKnownShaderAndRejectsUnknownName) { + const WebGPUShaderInfo& sigmoid = get_webgpu_shader_info("sigmoid"); + EXPECT_EQ(sigmoid.name, "sigmoid"); + EXPECT_NE(sigmoid.source, nullptr); + EXPECT_GT(sigmoid.workgroup_size_x, 0u); + EXPECT_THROW( + get_webgpu_shader_info("not_a_registered_shader"), std::runtime_error); +} + +TEST(WebGPUQ4RouteSignal, PreservesStaticAndRecordsBothDynamicSignals) { + WebGPUGraph static_graph; + static_graph.set_device(g_device); + build_q4_route_graph(static_graph, Q4RouteSignal::Static); + const auto static_dispatches = q4_dispatches(static_graph); + ASSERT_EQ(static_dispatches.size(), 1); + EXPECT_EQ(static_graph.num_dispatches(), 1); + EXPECT_NE(static_dispatches[0]->pipeline, nullptr); + EXPECT_NE(static_dispatches[0]->bind_group, nullptr); + EXPECT_FALSE(static_graph.has_dynamic_shapes()); + EXPECT_FALSE(static_graph.config().record_q4gsw_decode_route); + + WebGPUGraph legacy_graph; + legacy_graph.set_device(g_device); + build_q4_route_graph(legacy_graph, Q4RouteSignal::LegacyGraphMarker); + EXPECT_TRUE(legacy_graph.has_dynamic_shapes()); + EXPECT_FALSE(legacy_graph.config().record_q4gsw_decode_route); + EXPECT_EQ(legacy_graph.num_dispatches(), 3); + expect_dual_q4_topology(legacy_graph); + + WebGPUGraph explicit_graph; + explicit_graph.set_device(g_device); + build_q4_route_graph(explicit_graph, Q4RouteSignal::ExplicitOption); + EXPECT_FALSE(explicit_graph.has_dynamic_shapes()); + EXPECT_TRUE(explicit_graph.config().record_q4gsw_decode_route); + EXPECT_EQ(explicit_graph.num_dispatches(), 2); + expect_dual_q4_topology(explicit_graph); +} + +TEST(WebGPUComputeDispatch, PipelineKeyCanonicalizesConstants) { + WebGPUComputeDispatchDescriptor first; + first.shader_name = "sigmoid"; + first.entry_point = "main"; + first.constants = {{"beta", 2.0}, {"alpha", 1.0}}; + + WebGPUComputeDispatchDescriptor reordered = first; + reordered.constants = {{"alpha", 1.0}, {"beta", 2.0}}; + + EXPECT_EQ( + make_compute_pipeline_key(first), make_compute_pipeline_key(reordered)); +} + +TEST(WebGPUComputeDispatch, PipelineKeyTracksCompileIdentityOnly) { + WebGPUComputeDispatchDescriptor base; + base.shader_name = "sigmoid"; + base.entry_point = "main"; + base.constants = {{"wg_size", 64.0}}; + base.grid = {1u, 1u}; + base.bindings = {{reinterpret_cast(1), 0u, 16u}}; + + WebGPUComputeDispatchDescriptor runtime_change = base; + runtime_change.grid = {17u, 3u}; + runtime_change.bindings = {{reinterpret_cast(2), 128u, 4096u}}; + EXPECT_EQ( + make_compute_pipeline_key(base), + make_compute_pipeline_key(runtime_change)); + + WebGPUComputeDispatchDescriptor shader_change = base; + shader_change.shader_name = "binary_add"; + EXPECT_NE( + make_compute_pipeline_key(base), + make_compute_pipeline_key(shader_change)); + + WebGPUComputeDispatchDescriptor entry_change = base; + entry_change.entry_point = "alternate"; + EXPECT_NE( + make_compute_pipeline_key(base), make_compute_pipeline_key(entry_change)); + + WebGPUComputeDispatchDescriptor constant_change = base; + constant_change.constants = {{"wg_size", 128.0}}; + EXPECT_NE( + make_compute_pipeline_key(base), + make_compute_pipeline_key(constant_change)); +} + +TEST(WebGPUComputeDispatch, PipelineKeyRejectsInvalidConstants) { + WebGPUComputeDispatchDescriptor duplicate; + duplicate.shader_name = "sigmoid"; + duplicate.constants = {{"wg_size", 64.0}, {"wg_size", 128.0}}; + EXPECT_THROW(make_compute_pipeline_key(duplicate), std::runtime_error); + + WebGPUComputeDispatchDescriptor non_finite; + non_finite.shader_name = "sigmoid"; + non_finite.constants = {{"wg_size", std::numeric_limits::infinity()}}; + EXPECT_THROW(make_compute_pipeline_key(non_finite), std::runtime_error); +} + +TEST(WebGPUComputeDispatch, DescriptorRejectsInvalidBindings) { + WebGPUComputeDispatchDescriptor null_buffer; + null_buffer.shader_name = "sigmoid"; + null_buffer.bindings = {{nullptr, 0u, 16u}}; + EXPECT_THROW( + validate_compute_dispatch_descriptor(null_buffer), std::runtime_error); + + WebGPUComputeDispatchDescriptor zero_size; + zero_size.shader_name = "sigmoid"; + zero_size.bindings = {{reinterpret_cast(1), 0u, 0u}}; + EXPECT_THROW( + validate_compute_dispatch_descriptor(zero_size), std::runtime_error); + + WebGPUComputeDispatchDescriptor overflow; + overflow.shader_name = "sigmoid"; + overflow.bindings = { + {reinterpret_cast(1), + std::numeric_limits::max(), + 2u}}; + EXPECT_THROW( + validate_compute_dispatch_descriptor(overflow), std::runtime_error); +} + +TEST(WebGPUComputeDispatch, ReusesPipelineAndReleasesDawnObjects) { + constexpr size_t kNumElements = 64; + constexpr size_t kBufferBytes = kNumElements * sizeof(float); + + for (int iteration = 0; iteration < 32; iteration++) { + WGPUBuffer input = create_storage_buffer(kBufferBytes); + WGPUBuffer output = create_storage_buffer(kBufferBytes); + { + WebGPUGraph graph; + graph.set_device(g_device); + WGPUBuffer clamp_params = + graph.create_params_buffer(UnaryParams{kNumElements, -1.0f, 1.0f, 0}); + WGPUBuffer hardtanh_params = + graph.create_params_buffer(UnaryParams{kNumElements, -2.0f, 2.0f, 0}); + + WebGPUComputeDispatchDescriptor descriptor; + descriptor.shader_name = "clamp"; + descriptor.kernel_name = "clamp_test"; + descriptor.bindings = { + {input, 0u, kBufferBytes}, + {output, 0u, kBufferBytes}, + {clamp_params, 0u, sizeof(UnaryParams)}}; + descriptor.constants = {{"wg_size", 64.0}}; + descriptor.grid = {7u, 3u}; + + const size_t first = graph.add_compute_dispatch(descriptor); + descriptor.kernel_name = "hardtanh_test"; + descriptor.bindings[2].buffer = hardtanh_params; + const size_t second = graph.add_compute_dispatch(descriptor); + + const WebGPUMemoryStats stats = graph.memory_stats(); + EXPECT_EQ(stats.num_dispatches, 2); + EXPECT_EQ(stats.num_cached_shaders, 1); + EXPECT_EQ(stats.num_cached_pipelines, 1); + EXPECT_EQ(stats.uniform_buffer_bytes, 2 * sizeof(UnaryParams)); + EXPECT_EQ( + graph.dispatch_at(first).pipeline, + graph.dispatch_at(second).pipeline); + EXPECT_EQ(graph.dispatch_at(first).kernel_name, "clamp_test"); + EXPECT_EQ(graph.dispatch_at(second).kernel_name, "hardtanh_test"); + EXPECT_EQ(graph.dispatch_at(first).workgroup_count_x, 7u); + EXPECT_EQ(graph.dispatch_at(first).workgroup_count_y, 3u); + EXPECT_EQ(graph.dispatch_at(second).workgroup_count_x, 7u); + EXPECT_EQ(graph.dispatch_at(second).workgroup_count_y, 3u); + } + wgpuBufferRelease(output); + wgpuBufferRelease(input); + } +} + +TEST(WebGPUComputeDispatch, AlternateShaderReusesLayoutAndBindGroup) { + constexpr size_t kNumElements = 64; + constexpr size_t kBufferBytes = kNumElements * sizeof(float); + WGPUBuffer input = create_storage_buffer(kBufferBytes); + WGPUBuffer output = create_storage_buffer(kBufferBytes); + { + WebGPUGraph graph; + graph.set_device(g_device); + WGPUBuffer params = + graph.create_params_buffer(UnaryParams{kNumElements, 0.0f, 0.0f, 0}); + const std::vector bindings = { + {0, WGPUBufferBindingType_ReadOnlyStorage, input, kBufferBytes}, + {1, WGPUBufferBindingType_Storage, output, kBufferBytes}, + {2, WGPUBufferBindingType_Uniform, params, sizeof(UnaryParams)}, + }; + const WGPUConstantEntry wg_size = utils::make_wg_size_constant(64u); + { + utils::ComputePipelineBundle sigmoid = utils::make_compute_pipeline( + g_device, kSigmoidWGSL, bindings, &wg_size, 1u); + utils::ComputePipelineBundle relu = utils::make_compute_pipeline( + g_device, kReluWGSL, sigmoid, &wg_size, 1u); + + EXPECT_EQ(relu.bind_group, sigmoid.bind_group); + EXPECT_EQ(relu.bind_group_layout, nullptr); + EXPECT_EQ(relu.pipeline_layout, nullptr); + EXPECT_NE(relu.shader, nullptr); + EXPECT_NE(relu.pipeline, nullptr); + EXPECT_NE(relu.pipeline, sigmoid.pipeline); + + graph.add_dispatch( + {sigmoid.pipeline, sigmoid.bind_group, 1u, "sigmoid_test"}); + graph.add_dispatch({relu.pipeline, relu.bind_group, 1u, "relu_test"}); + } + } + wgpuBufferRelease(output); + wgpuBufferRelease(input); +} + +TEST(WebGPUComputeDispatch, RejectsBindingRangeBeyondDawnBuffer) { + WGPUBuffer buffer = create_storage_buffer(16u); + WebGPUComputeDispatchDescriptor descriptor; + descriptor.shader_name = "sigmoid"; + descriptor.bindings = {{buffer, 8u, 12u}}; + + EXPECT_THROW( + validate_compute_dispatch_descriptor(descriptor), std::runtime_error); + wgpuBufferRelease(buffer); +} + +constexpr int kResizeQ = 0; +constexpr int kResizeK = 1; +constexpr int kResizeUnrelated = 2; +constexpr int kResizeCascade = 3; +constexpr int kResizeOutput = 4; +constexpr int kResizeSymInt = 5; + +void build_resize_test_graph(WebGPUGraph& graph) { + namespace vk = vkgraph; + ::flatbuffers::FlatBufferBuilder fbb; + std::vector<::flatbuffers::Offset> values; + const std::vector dims = {8, 8}; + for (int mem_obj_id = 0; mem_obj_id < 5; mem_obj_id++) { + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::VkTensor, + vk::CreateVkTensorDirect( + fbb, + vk::VkDataType::FLOAT32, + &dims, + /*constant_id=*/-1, + mem_obj_id) + .Union())); + } + values.push_back(vk::CreateVkValue( + fbb, vk::GraphTypes::SymInt, vk::CreateSymInt(fbb, 0).Union())); + + const std::vector<::flatbuffers::Offset> chain; + const std::vector input_ids = { + kResizeQ, kResizeK, kResizeUnrelated}; + const std::vector output_ids = {kResizeOutput}; + const auto root = vk::CreateVkGraphDirect( + fbb, "0", &chain, &values, &input_ids, &output_ids); + vk::FinishVkGraphBuffer(fbb, root); + + graph.set_device(g_device); + graph.build(fbb.GetBufferPointer(), nullptr, 0, nullptr); +} + +WebGPUComputeDispatchDescriptor make_dynamic_test_descriptor( + WebGPUGraph& graph, + const char* kernel_name) { + const auto& input = graph.get_tensor(kResizeQ); + const auto& output = graph.get_tensor(kResizeOutput); + WGPUBuffer params = + graph.create_params_buffer(UnaryParams{64u, -1.0f, 1.0f, 0u}); + WebGPUComputeDispatchDescriptor descriptor; + descriptor.shader_name = "clamp"; + descriptor.kernel_name = kernel_name; + descriptor.bindings = { + {input.buffer, 0u, input.nbytes}, + {output.buffer, 0u, output.nbytes}, + {params, 0u, sizeof(UnaryParams)}}; + descriptor.constants = {{"wg_size", 64.0}}; + descriptor.grid = {99u, 97u}; + return descriptor; +} + +struct ResizeProbeContext { + int marker; + int* observed; + int* calls; +}; + +void record_resize_probe(WebGPUGraph&, const ResizeProbeContext& context) { + *context.observed = context.marker; + ++*context.calls; +} + +struct ThrowingResizeContext { + bool* fail; + int* calls; +}; + +void maybe_throw_resize(WebGPUGraph&, const ThrowingResizeContext& context) { + ++*context.calls; + if (*context.fail) { + throw std::runtime_error("resize hook failure"); + } +} + +struct GridPickerContext { + int tensor_id; + uint32_t x_bias; + uint32_t y_bias; + int* calls; + const bool* fail; +}; + +WebGPUDispatchGrid pick_tensor_grid( + const WebGPUGraph& graph, + const GridPickerContext& context) { + ++*context.calls; + if (context.fail != nullptr && *context.fail) { + throw std::runtime_error("dynamic grid picker failure"); + } + const auto& dims = graph.cur_dims(context.tensor_id); + return { + static_cast(dims.at(0)) + context.x_bias, + static_cast(dims.at(1)) + context.y_bias}; +} + +struct CascadeContext { + int source_id; + int output_id; +}; + +void resize_cascade_output(WebGPUGraph& graph, const CascadeContext& context) { + graph.set_cur_dims(context.output_id, graph.cur_dims(context.source_id)); +} + +void expect_dispatch_grid( + WebGPUGraph& graph, + size_t dispatch_index, + uint32_t x, + uint32_t y) { + const auto& dispatch = graph.dispatch_at(dispatch_index); + EXPECT_EQ(dispatch.workgroup_count_x, x); + EXPECT_EQ(dispatch.workgroup_count_y, y); +} + +TEST(WebGPUResizeHooks, TypedRegistrationOwnsContextCopies) { + WebGPUGraph graph; + build_resize_test_graph(graph); + int tensor_observed = 0; + int tensor_calls = 0; + int symint_observed = 0; + int symint_calls = 0; + using ResizeProbeFn = void (*)(WebGPUGraph&, const ResizeProbeContext&); + EXPECT_THROW( + graph.add_tensor_resize_hook( + kResizeQ, + static_cast(nullptr), + ResizeProbeContext{0, &tensor_observed, &tensor_calls}), + std::runtime_error); + EXPECT_THROW( + graph.add_tensor_resize_hook( + kResizeSymInt, + record_resize_probe, + ResizeProbeContext{0, &tensor_observed, &tensor_calls}), + std::runtime_error); + EXPECT_THROW( + graph.add_resize_hook( + kResizeQ, + record_resize_probe, + ResizeProbeContext{0, &symint_observed, &symint_calls}), + std::runtime_error); + { + ResizeProbeContext tensor_context = {17, &tensor_observed, &tensor_calls}; + ResizeProbeContext symint_context = {23, &symint_observed, &symint_calls}; + graph.add_tensor_resize_hook(kResizeQ, record_resize_probe, tensor_context); + graph.add_resize_hook(kResizeSymInt, record_resize_probe, symint_context); + tensor_context.marker = 101; + symint_context.marker = 103; + } + + graph.resize_input(kResizeQ, {7, 8}); + graph.propagate_resize(); + EXPECT_EQ(tensor_observed, 17); + EXPECT_EQ(tensor_calls, 1); + EXPECT_EQ(symint_observed, 0); + EXPECT_EQ(symint_calls, 0); + + graph.set_symint(kResizeSymInt, 9); + graph.propagate_resize(); + EXPECT_EQ(tensor_calls, 1); + EXPECT_EQ(symint_observed, 23); + EXPECT_EQ(symint_calls, 1); +} + +TEST(WebGPUResizeHooks, RestoresDirtyTriggerWhenHookThrows) { + WebGPUGraph graph; + build_resize_test_graph(graph); + bool hook_fails = true; + int hook_calls = 0; + int picker_calls = 0; + graph.add_tensor_resize_hook( + kResizeQ, + maybe_throw_resize, + ThrowingResizeContext{&hook_fails, &hook_calls}); + const size_t dispatch = graph.add_dynamic_compute_dispatch( + make_dynamic_test_descriptor(graph, "dynamic_hook_retry"), + kResizeQ, + pick_tensor_grid, + GridPickerContext{kResizeQ, 1u, 2u, &picker_calls, nullptr}); + picker_calls = 0; + + graph.resize_input(kResizeQ, {4, 3}); + EXPECT_THROW(graph.propagate_resize(), std::runtime_error); + EXPECT_EQ(hook_calls, 1); + EXPECT_EQ(picker_calls, 0); + expect_dispatch_grid(graph, dispatch, 9u, 10u); + + hook_fails = false; + EXPECT_NO_THROW(graph.propagate_resize()); + EXPECT_EQ(hook_calls, 2); + EXPECT_EQ(picker_calls, 1); + expect_dispatch_grid(graph, dispatch, 5u, 5u); +} + +TEST(WebGPUDynamicDispatch, InitializesAndIsolatesTriggeredGrids) { + WebGPUGraph graph; + build_resize_test_graph(graph); + int q_calls = 0; + int k_calls = 0; + GridPickerContext q_context = {kResizeQ, 1u, 2u, &q_calls, nullptr}; + const GridPickerContext k_context = {kResizeK, 3u, 4u, &k_calls, nullptr}; + const size_t q_dispatch = graph.add_dynamic_compute_dispatch( + make_dynamic_test_descriptor(graph, "dynamic_q"), + kResizeQ, + pick_tensor_grid, + q_context); + const size_t k_dispatch = graph.add_dynamic_compute_dispatch( + make_dynamic_test_descriptor(graph, "dynamic_k"), + kResizeK, + pick_tensor_grid, + k_context); + q_context.x_bias = 101u; + q_context.y_bias = 103u; + expect_dispatch_grid(graph, q_dispatch, 9u, 10u); + expect_dispatch_grid(graph, k_dispatch, 11u, 12u); + EXPECT_EQ(q_calls, 1); + EXPECT_EQ(k_calls, 1); + q_calls = 0; + k_calls = 0; + + graph.resize_input(kResizeUnrelated, {7, 7}); + graph.propagate_resize(); + EXPECT_EQ(q_calls, 0); + EXPECT_EQ(k_calls, 0); + + graph.resize_input(kResizeQ, {4, 3}); + graph.propagate_resize(); + expect_dispatch_grid(graph, q_dispatch, 5u, 5u); + expect_dispatch_grid(graph, k_dispatch, 11u, 12u); + EXPECT_EQ(q_calls, 1); + EXPECT_EQ(k_calls, 0); + + graph.resize_input(kResizeQ, {2, 5}); + graph.propagate_resize(); + expect_dispatch_grid(graph, q_dispatch, 3u, 7u); + EXPECT_EQ(q_calls, 2); + graph.resize_input(kResizeQ, {2, 5}); + graph.propagate_resize(); + EXPECT_EQ(q_calls, 2); + + graph.resize_input(kResizeK, {6, 1}); + graph.propagate_resize(); + expect_dispatch_grid(graph, k_dispatch, 9u, 5u); + EXPECT_EQ(k_calls, 1); +} + +TEST(WebGPUDynamicDispatch, HandlesCascadesAndStagesPickerFailures) { + { + WebGPUGraph graph; + build_resize_test_graph(graph); + int same_pass_calls = 0; + int cascade_pass_calls = 0; + graph.add_tensor_resize_hook( + kResizeQ, + resize_cascade_output, + CascadeContext{kResizeQ, kResizeCascade}); + const size_t same_pass_dispatch = graph.add_dynamic_compute_dispatch( + make_dynamic_test_descriptor(graph, "dynamic_same_pass"), + kResizeQ, + pick_tensor_grid, + GridPickerContext{kResizeCascade, 5u, 7u, &same_pass_calls, nullptr}); + const size_t cascade_pass_dispatch = graph.add_dynamic_compute_dispatch( + make_dynamic_test_descriptor(graph, "dynamic_cascade_pass"), + kResizeCascade, + pick_tensor_grid, + GridPickerContext{ + kResizeCascade, 9u, 11u, &cascade_pass_calls, nullptr}); + same_pass_calls = 0; + cascade_pass_calls = 0; + graph.resize_input(kResizeQ, {4, 6}); + graph.propagate_resize(); + EXPECT_EQ(same_pass_calls, 1); + EXPECT_EQ(cascade_pass_calls, 1); + expect_dispatch_grid(graph, same_pass_dispatch, 9u, 13u); + expect_dispatch_grid(graph, cascade_pass_dispatch, 13u, 17u); + } + + WebGPUGraph graph; + build_resize_test_graph(graph); + int first_calls = 0; + int second_calls = 0; + bool second_fails = false; + const size_t first = graph.add_dynamic_compute_dispatch( + make_dynamic_test_descriptor(graph, "dynamic_first"), + kResizeQ, + pick_tensor_grid, + GridPickerContext{kResizeQ, 1u, 2u, &first_calls, nullptr}); + const size_t second = graph.add_dynamic_compute_dispatch( + make_dynamic_test_descriptor(graph, "dynamic_second"), + kResizeQ, + pick_tensor_grid, + GridPickerContext{kResizeQ, 3u, 4u, &second_calls, &second_fails}); + expect_dispatch_grid(graph, first, 9u, 10u); + expect_dispatch_grid(graph, second, 11u, 12u); + second_fails = true; + graph.resize_input(kResizeQ, {4, 3}); + EXPECT_THROW(graph.propagate_resize(), std::runtime_error); + expect_dispatch_grid(graph, first, 9u, 10u); + expect_dispatch_grid(graph, second, 11u, 12u); + second_fails = false; + EXPECT_NO_THROW(graph.propagate_resize()); + expect_dispatch_grid(graph, first, 5u, 5u); + expect_dispatch_grid(graph, second, 7u, 7u); +} + +TEST(WebGPUDynamicDispatch, RejectsRouteOverlapWithoutPoisoningRegistry) { + WebGPUGraph graph; + build_resize_test_graph(graph); + int calls = 0; + const size_t dispatches_before_invalid_trigger = graph.num_dispatches(); + EXPECT_THROW( + graph.add_dynamic_compute_dispatch( + make_dynamic_test_descriptor(graph, "dynamic_negative_trigger"), + -1, + pick_tensor_grid, + GridPickerContext{kResizeQ, 0u, 0u, &calls, nullptr}), + std::runtime_error); + EXPECT_THROW( + graph.add_dynamic_compute_dispatch( + make_dynamic_test_descriptor(graph, "dynamic_oob_trigger"), + graph.num_values(), + pick_tensor_grid, + GridPickerContext{kResizeQ, 0u, 0u, &calls, nullptr}), + std::runtime_error); + EXPECT_THROW( + graph.add_dynamic_compute_dispatch( + make_dynamic_test_descriptor(graph, "dynamic_symint_trigger"), + kResizeSymInt, + pick_tensor_grid, + GridPickerContext{kResizeQ, 0u, 0u, &calls, nullptr}), + std::runtime_error); + EXPECT_EQ(calls, 0); + EXPECT_EQ(graph.num_dispatches(), dispatches_before_invalid_trigger); + + auto pick_zero_grid = [](const WebGPUGraph&, const WebGPUDispatchGrid& grid) { + return grid; + }; + EXPECT_THROW( + graph.add_dynamic_compute_dispatch( + make_dynamic_test_descriptor(graph, "dynamic_zero_x"), + kResizeQ, + +pick_zero_grid, + WebGPUDispatchGrid{0u, 1u}), + std::runtime_error); + EXPECT_THROW( + graph.add_dynamic_compute_dispatch( + make_dynamic_test_descriptor(graph, "dynamic_zero_y"), + kResizeQ, + +pick_zero_grid, + WebGPUDispatchGrid{1u, 0u}), + std::runtime_error); + EXPECT_EQ(graph.num_dispatches(), dispatches_before_invalid_trigger); + + bool initial_fails = true; + const size_t dispatches_before_failure = graph.num_dispatches(); + EXPECT_THROW( + graph.add_dynamic_compute_dispatch( + make_dynamic_test_descriptor(graph, "dynamic_initial_failure"), + kResizeQ, + pick_tensor_grid, + GridPickerContext{kResizeQ, 0u, 0u, &calls, &initial_fails}), + std::runtime_error); + EXPECT_EQ(graph.num_dispatches(), dispatches_before_failure); + + const size_t dynamic = graph.add_dynamic_compute_dispatch( + make_dynamic_test_descriptor(graph, "dynamic_route_guard"), + kResizeQ, + pick_tensor_grid, + GridPickerContext{kResizeQ, 0u, 0u, &calls, nullptr}); + ASSERT_EQ(dynamic, 0u); + graph.add_dispatch(WebGPUDispatch{}); + graph.add_dispatch(WebGPUDispatch{}); + + EXPECT_THROW( + graph.register_dispatch_route_group({{0, 1}, {1, 2}}), + std::runtime_error); + calls = 0; + graph.resize_input(kResizeQ, {7, 6}); + graph.propagate_resize(); + EXPECT_EQ(calls, 1); + expect_dispatch_grid(graph, dynamic, 7u, 6u); + const size_t group = graph.register_dispatch_route_group({{1, 2}, {2, 3}}); + EXPECT_EQ(group, 0u); + graph.select_dispatch_route(group, 1, {{13u, 17u}}); + expect_dispatch_grid(graph, 1u, 0u, 0u); + expect_dispatch_grid(graph, 2u, 13u, 17u); +} + +struct InvalidRopeGraphCase { + const char* name; + std::vector xq_dims; + std::vector xk_dims; + std::vector cos_dims; + std::vector sin_dims; + std::vector xq_out_dims; + std::vector xk_out_dims; + vkgraph::VkDataType xq_dtype; + const char* expected_error; +}; + +void expect_invalid_rope_graph(const InvalidRopeGraphCase& test_case) { + namespace vk = vkgraph; + ::flatbuffers::FlatBufferBuilder fbb; + std::vector<::flatbuffers::Offset> values; + auto add_tensor = [&](vk::VkDataType dtype, + const std::vector& dims, + int mem_obj_id) { + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::VkTensor, + vk::CreateVkTensorDirect( + fbb, dtype, &dims, /*constant_id=*/-1, mem_obj_id) + .Union())); + }; + add_tensor(test_case.xq_dtype, test_case.xq_dims, 0); + add_tensor(vk::VkDataType::FLOAT32, test_case.xk_dims, 1); + add_tensor(vk::VkDataType::FLOAT32, test_case.cos_dims, 2); + add_tensor(vk::VkDataType::FLOAT32, test_case.sin_dims, 3); + add_tensor(vk::VkDataType::FLOAT32, test_case.xq_out_dims, 4); + add_tensor(vk::VkDataType::FLOAT32, test_case.xk_out_dims, 5); + std::vector output_value_ids = {4, 5}; + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::ValueList, + vk::CreateValueListDirect(fbb, &output_value_ids).Union())); + + std::vector args = {0, 1, 2, 3, 6}; + std::vector<::flatbuffers::Offset> chain; + chain.push_back(vk::CreateOperatorCallDirect( + fbb, 0, "et_vk.apply_rotary_emb.default", &args)); + std::vector input_ids = {0, 1, 2, 3}; + std::vector output_ids = {4, 5}; + const auto root = vk::CreateVkGraphDirect( + fbb, "0", &chain, &values, &input_ids, &output_ids); + vk::FinishVkGraphBuffer(fbb, root); + + WebGPUGraph graph; + std::string error; + try { + graph.build(fbb.GetBufferPointer(), nullptr, 0, nullptr); + } catch (const std::exception& exception) { + error = exception.what(); + } + EXPECT_FALSE(error.empty()) << test_case.name << " unexpectedly built"; + EXPECT_EQ(error, test_case.expected_error) + << test_case.name << " rejected for the wrong reason"; + const WebGPUMemoryStats stats = graph.memory_stats(); + EXPECT_EQ(stats.num_dispatches, 0) << test_case.name; + EXPECT_EQ(stats.uniform_buffer_bytes, 0u) << test_case.name; + EXPECT_EQ(stats.num_cached_shaders, 0) << test_case.name; + EXPECT_EQ(stats.num_cached_pipelines, 0) << test_case.name; +} + +TEST(WebGPURopeValidation, RejectsMalformedGraphsBeforeDispatchAllocation) { + ASSERT_TRUE( + webgpu_operator_registry().has_op("et_vk.apply_rotary_emb.default")); + const std::vector xq = {1, 2, 2, 4}; + const std::vector xk = {1, 2, 1, 4}; + const std::vector freqs = {2, 2}; + const InvalidRopeGraphCase cases[] = { + {"query rank", + {2, 4}, + xk, + freqs, + freqs, + {2, 4}, + xk, + vkgraph::VkDataType::FLOAT32, + "WebGPU apply_rotary_emb: malformed dims"}, + {"sequence mismatch", + xq, + {1, 3, 1, 4}, + freqs, + freqs, + xq, + {1, 3, 1, 4}, + vkgraph::VkDataType::FLOAT32, + "WebGPU apply_rotary_emb: xq/xk head_dim and seq must match"}, + {"head dimension mismatch", + xq, + {1, 2, 1, 6}, + freqs, + freqs, + xq, + {1, 2, 1, 6}, + vkgraph::VkDataType::FLOAT32, + "WebGPU apply_rotary_emb: xq/xk head_dim and seq must match"}, + {"frequency width mismatch", + xq, + xk, + {2, 3}, + {2, 3}, + xq, + xk, + vkgraph::VkDataType::FLOAT32, + "WebGPU apply_rotary_emb: head_dim != 2 * freqs_cos last dim"}, + {"cosine/sine shape mismatch", + xq, + xk, + freqs, + {2, 1}, + xq, + xk, + vkgraph::VkDataType::FLOAT32, + "WebGPU apply_rotary_emb: freqs_cos and freqs_sin shapes differ"}, + {"query byte size mismatch", + xq, + xk, + freqs, + freqs, + xq, + xk, + vkgraph::VkDataType::INT64, + "WebGPU apply_rotary_emb: dtype/byte-size mismatch (all fp32) or " + "freqs shape != [seq, head_dim/2]"}, + }; + for (const InvalidRopeGraphCase& test_case : cases) { + SCOPED_TRACE(test_case.name); + expect_invalid_rope_graph(test_case); + } +} + +TEST(WebGPUExecution, FullySuppressedPlanPerformsNoQueueSubmission) { + WebGPUGraph graph; + const WebGPUExecutionPlan plan; + + EXPECT_EQ(graph.execute(plan), 0u); +} + +TEST(WebGPUExecution, RejectsPlanOutputCountMismatch) { + WebGPUGraph graph; + WebGPUExecutionPlan plan; + plan.copy_outputs = {true}; + std::vector outputs; + + EXPECT_THROW(graph.execute(plan), std::runtime_error); + EXPECT_THROW(graph.copy_outputs(outputs, plan), std::runtime_error); +} + +TEST(WebGPUExecution, RejectsPlanDispatchOutOfRange) { + WebGPUGraph graph; + WebGPUExecutionPlan plan; + plan.dispatch_chunks = {{0u}}; + + EXPECT_THROW(graph.execute(plan), std::runtime_error); +} + +} // namespace +} // namespace executorch::backends::webgpu + +int main(int argc, char** argv) { + ::testing::InitGoogleTest(&argc, argv); + + executorch::backends::webgpu::WebGPUContext context; + try { + context = executorch::backends::webgpu::create_webgpu_context(); + } catch (const std::exception& error) { + std::printf("SKIP: %s\n", error.what()); + return 0; + } + executorch::backends::webgpu::set_default_webgpu_context(&context); + executorch::backends::webgpu::g_device = context.device; + + const int result = RUN_ALL_TESTS(); + executorch::backends::webgpu::set_default_webgpu_context(nullptr); + executorch::backends::webgpu::destroy_webgpu_context(context); + return result; +} diff --git a/backends/webgpu/test/native/test_dispatch_2d.cpp b/backends/webgpu/test/native/test_dispatch_2d.cpp index 582caa3c98e..7c5a5e141f3 100644 --- a/backends/webgpu/test/native/test_dispatch_2d.cpp +++ b/backends/webgpu/test/native/test_dispatch_2d.cpp @@ -9,13 +9,21 @@ // Device-free unit test for the pure 2D workgroup-count fold that lifts the // 65535 per-dim dispatch cap. Exercises the fold arithmetic only — no GPU. +#include +#include #include #include #include #include +#include +#include +using executorch::backends::webgpu::WebGPUDispatch; +using executorch::backends::webgpu::WebGPUGraph; +using executorch::backends::webgpu::utils::DispatchRange; +using executorch::backends::webgpu::utils::DispatchRouteRegistry; using executorch::backends::webgpu::utils::fold_workgroup_count_2d; using executorch::backends::webgpu::utils::WgCount; @@ -57,4 +65,257 @@ TEST(DispatchFold, ThrowsWhenNeeds3rdDimension) { EXPECT_ANY_THROW(fold_workgroup_count_2d(kMax * kMax + 1u, kMax, "test")); } +void expect_grid(const WgCount& grid, uint32_t x, uint32_t y) { + EXPECT_EQ(grid.x, x); + EXPECT_EQ(grid.y, y); +} + +void expect_grids_equal( + const std::vector& actual, + const std::vector& expected) { + ASSERT_EQ(actual.size(), expected.size()); + for (size_t i = 0; i < actual.size(); i++) { + EXPECT_EQ(actual[i].x, expected[i].x) << "index=" << i; + EXPECT_EQ(actual[i].y, expected[i].y) << "index=" << i; + } +} + +TEST(DispatchRoute, SwitchesUnequalRangesAndRestoresBothDimensions) { + std::vector grids(6, {101, 103}); + const std::vector compute(6, true); + DispatchRouteRegistry registry; + const size_t group = registry.register_group( + grids.size(), {{1, 3}, {4, 5}}, [&](size_t i) { return compute[i]; }); + auto set_grid = [&](size_t i, WgCount grid) { grids[i] = grid; }; + + registry.select(group, 0, {{3, 7}, {5, 11}}, set_grid); + expect_grid(grids[0], 101, 103); + expect_grid(grids[1], 3, 7); + expect_grid(grids[2], 5, 11); + expect_grid(grids[3], 101, 103); + expect_grid(grids[4], 0, 0); + expect_grid(grids[5], 101, 103); + + registry.select(group, 1, {{13, 17}}, set_grid); + expect_grid(grids[1], 0, 0); + expect_grid(grids[2], 0, 0); + expect_grid(grids[4], 13, 17); + + registry.select(group, 0, {{19, 23}, {29, 31}}, set_grid); + expect_grid(grids[1], 19, 23); + expect_grid(grids[2], 29, 31); + expect_grid(grids[4], 0, 0); +} + +TEST(DispatchRoute, InvalidSelectionDoesNotMutateGrids) { + std::vector grids = {{1, 2}, {3, 4}, {5, 6}}; + const std::vector original = grids; + DispatchRouteRegistry registry; + const size_t group = registry.register_group( + grids.size(), {{0, 2}, {2, 3}}, [](size_t) { return true; }); + auto set_grid = [&](size_t i, WgCount grid) { grids[i] = grid; }; + + EXPECT_ANY_THROW(registry.select(group, 2, {{7, 8}}, set_grid)); + expect_grids_equal(grids, original); + EXPECT_ANY_THROW(registry.select(group, 0, {{7, 8}}, set_grid)); + expect_grids_equal(grids, original); + EXPECT_ANY_THROW(registry.select(group, 0, {{0, 8}, {9, 10}}, set_grid)); + expect_grids_equal(grids, original); + EXPECT_ANY_THROW(registry.select(group, 0, {{7, 0}, {9, 10}}, set_grid)); + expect_grids_equal(grids, original); +} + +TEST(DispatchRoute, RejectsInvalidAndMultiplyOwnedRanges) { + const auto all_compute = [](size_t) { return true; }; + DispatchRouteRegistry registry; + // Each pairs the range under test with a valid second range so the group + // clears the size < 2 short-circuit and the range-validation loop is actually + // exercised: {2,1} inverted, {1,1} empty, {0,5} end past dispatch_count. + EXPECT_ANY_THROW(registry.register_group(4, {{2, 1}, {3, 4}}, all_compute)); + EXPECT_ANY_THROW(registry.register_group(4, {{1, 1}, {3, 4}}, all_compute)); + EXPECT_ANY_THROW(registry.register_group(4, {{0, 5}, {3, 4}}, all_compute)); + EXPECT_ANY_THROW(registry.register_group(4, {{0, 2}, {1, 3}}, all_compute)); + + const size_t first = + registry.register_group(4, {{0, 1}, {1, 2}}, all_compute); + EXPECT_EQ(first, 0); + EXPECT_ANY_THROW(registry.register_group(4, {{1, 3}}, all_compute)); + + DispatchRouteRegistry copy_registry; + const std::vector compute = {true, false, true}; + EXPECT_ANY_THROW(copy_registry.register_group( + compute.size(), {{0, 1}, {1, 2}}, [&](size_t i) { return compute[i]; })); +} + +TEST(DispatchRoute, GraphRejectsCopyAndCrossGroupOwnership) { + WebGPUGraph graph; + for (size_t i = 0; i < 6; i++) { + graph.add_dispatch(WebGPUDispatch{}); + } + graph.add_buffer_copy(nullptr, nullptr, 0); + + const size_t group = graph.register_dispatch_route_group({{0, 1}, {1, 2}}); + EXPECT_EQ(group, 0); + EXPECT_ANY_THROW(graph.register_dispatch_route_group({{1, 2}, {2, 3}})); + EXPECT_ANY_THROW(graph.register_dispatch_route_group({{5, 6}, {6, 7}})); + + const size_t second = graph.register_dispatch_route_group({{2, 4}, {4, 5}}); + EXPECT_EQ(second, 1); + + graph.select_dispatch_route(group, 0, {{7, 11}}); + graph.select_dispatch_route(second, 0, {{13, 17}, {19, 23}}); + expect_grid( + {graph.dispatch_at(0).workgroup_count_x, + graph.dispatch_at(0).workgroup_count_y}, + 7, + 11); + expect_grid( + {graph.dispatch_at(1).workgroup_count_x, + graph.dispatch_at(1).workgroup_count_y}, + 0, + 0); + expect_grid( + {graph.dispatch_at(2).workgroup_count_x, + graph.dispatch_at(2).workgroup_count_y}, + 13, + 17); + expect_grid( + {graph.dispatch_at(3).workgroup_count_x, + graph.dispatch_at(3).workgroup_count_y}, + 19, + 23); + expect_grid( + {graph.dispatch_at(4).workgroup_count_x, + graph.dispatch_at(4).workgroup_count_y}, + 0, + 0); + + graph.select_dispatch_route(group, 1, {{29, 31}}); + expect_grid( + {graph.dispatch_at(2).workgroup_count_x, + graph.dispatch_at(2).workgroup_count_y}, + 13, + 17); +} + +TEST(DispatchRoute, ExecuteRejectsHalfZeroGrid) { + WebGPUGraph graph; + WebGPUDispatch dispatch; + dispatch.workgroup_count_x = 0; + dispatch.workgroup_count_y = 1; + graph.add_dispatch(dispatch); + EXPECT_ANY_THROW(graph.make_execution_plan({})); +} + +TEST(DispatchRoute, RecordsEligibleQ4AndDynamicSdpaAlternates) { + using executorch::backends::webgpu::utils::should_record_q4gsw_dual_route; + using executorch::backends::webgpu::utils::should_record_sdpa_dual_route; + + EXPECT_FALSE(should_record_q4gsw_dual_route(1, true, true, true)); + EXPECT_FALSE(should_record_q4gsw_dual_route(32, false, true, true)); + EXPECT_FALSE(should_record_q4gsw_dual_route(32, true, false, false)); + EXPECT_TRUE(should_record_q4gsw_dual_route(32, true, true, false)); + EXPECT_TRUE(should_record_q4gsw_dual_route(32, true, false, true)); + + EXPECT_FALSE(should_record_sdpa_dual_route(true, false, false)); + EXPECT_FALSE(should_record_sdpa_dual_route(false, true, true)); + EXPECT_TRUE(should_record_sdpa_dual_route(true, true, false)); + EXPECT_TRUE(should_record_sdpa_dual_route(true, false, true)); +} + +TEST(WebGPUGraphConfig, ParsesExactBooleanCompileOption) { + using executorch::backends::webgpu::parse_webgpu_graph_config; + using executorch::runtime::CompileSpec; + + auto absent = parse_webgpu_graph_config({}); + ASSERT_TRUE(absent.ok()); + EXPECT_FALSE(absent->record_q4gsw_decode_route); + EXPECT_FALSE(absent->f16_kv_cache); + EXPECT_FALSE(absent->f16_accumulate_gemm); + EXPECT_EQ(absent->sdpa_query_tile, 0); + + uint8_t false_value = 0; + CompileSpec false_spec = { + "webgpu_record_q4gsw_decode_route", {&false_value, 1}}; + auto parsed_false = parse_webgpu_graph_config(false_spec); + ASSERT_TRUE(parsed_false.ok()); + EXPECT_FALSE(parsed_false->record_q4gsw_decode_route); + + uint8_t true_value = 1; + CompileSpec true_spec = { + "webgpu_record_q4gsw_decode_route", {&true_value, 1}}; + auto parsed_true = parse_webgpu_graph_config(true_spec); + ASSERT_TRUE(parsed_true.ok()); + EXPECT_TRUE(parsed_true->record_q4gsw_decode_route); + + CompileSpec unknown_spec = {"unknown_webgpu_option", {nullptr, 0}}; + auto parsed_unknown = parse_webgpu_graph_config(unknown_spec); + ASSERT_TRUE(parsed_unknown.ok()); + EXPECT_FALSE(parsed_unknown->record_q4gsw_decode_route); +} + +TEST(WebGPUGraphConfig, RejectsMalformedRouteCompileOption) { + using executorch::backends::webgpu::parse_webgpu_graph_config; + using executorch::runtime::CompileSpec; + using executorch::runtime::Error; + + uint8_t values[2] = {0, 1}; + CompileSpec empty = {"webgpu_record_q4gsw_decode_route", {values, 0}}; + auto parsed_empty = parse_webgpu_graph_config(empty); + ASSERT_FALSE(parsed_empty.ok()); + EXPECT_EQ(parsed_empty.error(), Error::DelegateInvalidCompatibility); + + CompileSpec oversized = {"webgpu_record_q4gsw_decode_route", {values, 2}}; + auto parsed_oversized = parse_webgpu_graph_config(oversized); + ASSERT_FALSE(parsed_oversized.ok()); + EXPECT_EQ(parsed_oversized.error(), Error::DelegateInvalidCompatibility); + + CompileSpec null_value = {"webgpu_record_q4gsw_decode_route", {nullptr, 1}}; + auto parsed_null = parse_webgpu_graph_config(null_value); + ASSERT_FALSE(parsed_null.ok()); + EXPECT_EQ(parsed_null.error(), Error::DelegateInvalidCompatibility); +} + +TEST(DispatchRoute, Bk64RequiresExactLlamaShapeAndCapabilities) { + using executorch::backends::webgpu::utils::is_q4gsw_bk64_eligible; + + constexpr uint32_t kRequiredInvocations = 256u; + constexpr uint32_t kRequiredStorageBytes = 16384u; + for (const auto& shape : std::vector>{ + {2048u, 8192u}, {8192u, 2048u}, {2048u, 2048u}}) { + EXPECT_TRUE(is_q4gsw_bk64_eligible( + shape.first, + shape.second, + 64u, + false, + true, + kRequiredInvocations, + kRequiredStorageBytes)); + } + + EXPECT_FALSE(is_q4gsw_bk64_eligible( + 2048u, 512u, 64u, false, true, 256u, kRequiredStorageBytes)); + EXPECT_FALSE(is_q4gsw_bk64_eligible( + 2048u, 2048u, 32u, false, true, 256u, kRequiredStorageBytes)); + EXPECT_FALSE(is_q4gsw_bk64_eligible( + 2048u, 2048u, 64u, true, true, 256u, kRequiredStorageBytes)); + EXPECT_FALSE(is_q4gsw_bk64_eligible( + 2048u, 2048u, 64u, false, false, 256u, kRequiredStorageBytes)); + EXPECT_FALSE(is_q4gsw_bk64_eligible( + 2048u, 2048u, 64u, false, true, 255u, kRequiredStorageBytes)); + EXPECT_FALSE(is_q4gsw_bk64_eligible( + 2048u, 2048u, 64u, false, true, 256u, kRequiredStorageBytes - 1u)); +} + +TEST(DispatchRoute, Bk64SelectsOnlyAcceptedExactLiveRows) { + using executorch::backends::webgpu::utils::is_q4gsw_bk64_live_m; + + EXPECT_TRUE(is_q4gsw_bk64_live_m(128u)); + EXPECT_TRUE(is_q4gsw_bk64_live_m(508u)); + EXPECT_TRUE(is_q4gsw_bk64_live_m(512u)); + for (uint32_t m : {1u, 127u, 129u, 507u, 509u, 511u, 513u}) { + EXPECT_FALSE(is_q4gsw_bk64_live_m(m)) << "M=" << m; + } +} + } // namespace diff --git a/backends/webgpu/test/native/test_dynamic_shape.cpp b/backends/webgpu/test/native/test_dynamic_shape.cpp index 3752e4ee53f..7b3c1d6b0a7 100644 --- a/backends/webgpu/test/native/test_dynamic_shape.cpp +++ b/backends/webgpu/test/native/test_dynamic_shape.cpp @@ -24,17 +24,22 @@ // /tmp/dynamic_shape. #include +#include #include #include +#include +#include #include +#include #include #include #include #include #include #include +#include #include #include @@ -49,6 +54,26 @@ constexpr int kHidden = 64; // Artifacts directory; set from env/argv in main() before RUN_ALL_TESTS(). std::string g_dir; // NOLINT(cppcoreguidelines-avoid-non-const-global-variables) +#ifdef WGPU_BACKEND_ENABLE_PROFILING +std::vector current_profile_names() { + const auto* context = get_default_webgpu_context(); + if (context == nullptr || context->querypool == nullptr) { + return {}; + } + std::vector names; + for (const auto& duration : context->querypool->results()) { + names.push_back(duration.kernel_name); + } + return names; +} + +bool contains_name( + const std::vector& names, + const std::string& expected) { + return std::find(names.begin(), names.end(), expected) != names.end(); +} +#endif + std::vector read_bin(const std::string& path) { std::ifstream f(path, std::ios::binary | std::ios::ate); if (!f) { @@ -105,16 +130,27 @@ void check_s(Module& m, const std::string& prefix, int s) { // Dynamic quantized linear: input [M, kLinK] -> output [M, n]. kLinN is the // register-tiled/bicol config; kLinNShmem (N>=2048) routes to the shmem GEMM. constexpr int kLinK = 64; +constexpr int kLinAltK = 72; constexpr int kLinN = 128; constexpr int kLinNShmem = 2048; +constexpr int kFp32LinearN = 32; // Run at [m_rows, kLinK] on an already-loaded module (so it can be // reused across M without a fresh load), and compare to the golden. -void run_linear(Module& m, int m_rows, const char* prefix, int n) { +void run_linear( + Module& m, + int m_rows, + const char* prefix, + int n, + int k = kLinK, + float atol = 5e-3f, + float rtol = 0.0f, + float nrmse_limit = -1.0f, + float tail_nrmse_limit = -1.0f) { const std::string base = g_dir + "/" + prefix + ".S" + std::to_string(m_rows); auto input = read_bin(base + ".input.bin"); auto golden = read_bin(base + ".golden.bin"); ASSERT_FALSE(input.empty()) << "missing " << prefix << ".S" << m_rows; - auto t = make_tensor_ptr({m_rows, kLinK}, std::move(input)); + auto t = make_tensor_ptr({m_rows, k}, std::move(input)); auto r = m.forward({EValue(t)}); ASSERT_TRUE(r.ok() && !r.get().empty() && r.get()[0].isTensor()) << prefix << " M=" << m_rows << " forward failed"; @@ -124,9 +160,60 @@ void run_linear(Module& m, int m_rows, const char* prefix, int n) { << prefix << " M=" << m_rows << " output numel mismatch"; std::vector got( out.const_data_ptr(), out.const_data_ptr() + numel); - const float e = max_err(got, golden); - // 4-bit quant: looser tol (the kernel mirrors the dequant-matmul reference). - EXPECT_LT(e, 5e-3f) << prefix << " M=" << m_rows << " max_err=" << e; + ASSERT_EQ(got.size(), golden.size()); + float max_abs = 0.0f; + float max_rel = 0.0f; + double error_sq_sum = 0.0; + double golden_sq_sum = 0.0; + bool within_tolerance = true; + for (size_t i = 0; i < got.size(); ++i) { + ASSERT_TRUE(std::isfinite(got[i])) + << prefix << " M=" << m_rows << " i=" << i; + ASSERT_TRUE(std::isfinite(golden[i])) + << prefix << " M=" << m_rows << " golden i=" << i; + const float abs_err = std::fabs(got[i] - golden[i]); + const float rel_err = abs_err / std::fmax(std::fabs(golden[i]), 1e-6f); + max_abs = std::fmax(max_abs, abs_err); + max_rel = std::fmax(max_rel, rel_err); + if (abs_err > atol && rel_err > rtol) { + within_tolerance = false; + } + error_sq_sum += static_cast(abs_err) * abs_err; + golden_sq_sum += static_cast(golden[i]) * golden[i]; + } + EXPECT_TRUE(within_tolerance) + << prefix << " M=" << m_rows << " max_abs=" << max_abs + << " max_rel=" << max_rel << " tolerances=" << atol << "/" << rtol; + if (nrmse_limit > 0.0f) { + ASSERT_GT(golden_sq_sum, 0.0) + << prefix << " M=" << m_rows << " zero golden norm (NRMSE undefined)"; + const double nrmse = std::sqrt(error_sq_sum / golden_sq_sum); + EXPECT_LT(nrmse, nrmse_limit) + << prefix << " M=" << m_rows << " full-output NRMSE"; + std::printf( + "%s M=%d max_abs=%g max_rel=%g nrmse=%g\n", + prefix, + m_rows, + max_abs, + max_rel, + nrmse); + } + if (tail_nrmse_limit > 0.0f) { + double tail_error_sq_sum = 0.0; + double tail_golden_sq_sum = 0.0; + const size_t tail_begin = static_cast(m_rows - 1) * n; + for (size_t i = tail_begin; i < got.size(); ++i) { + const double error = static_cast(got[i]) - golden[i]; + tail_error_sq_sum += error * error; + tail_golden_sq_sum += static_cast(golden[i]) * golden[i]; + } + ASSERT_GT(tail_golden_sq_sum, 0.0) + << prefix << " M=" << m_rows << " zero final-row golden norm"; + const double tail_nrmse = std::sqrt(tail_error_sq_sum / tail_golden_sq_sum); + EXPECT_LT(tail_nrmse, tail_nrmse_limit) + << prefix << " M=" << m_rows << " final-row NRMSE"; + std::printf("%s M=%d final_row_nrmse=%g\n", prefix, m_rows, tail_nrmse); + } } void check_linear(int m_rows) { @@ -138,16 +225,293 @@ void check_linear(int m_rows) { void check_linear_shmem(int m_rows) { Module m(g_dir + "/dyn_linear_shmem.pte"); ASSERT_EQ(m.load_forward(), Error::Ok) << "load dyn_linear_shmem.pte"; - run_linear(m, m_rows, "dyn_linear_shmem", kLinNShmem); + run_linear(m, m_rows, "dyn_linear_shmem", kLinNShmem, kLinAltK); +} + +void check_linear_tiled(int m_rows) { + Module m(g_dir + "/dyn_linear_tiled.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "load dyn_linear_tiled.pte"; + run_linear(m, m_rows, "dyn_linear_tiled", kLinN, kLinAltK); +} + +void check_fp32_linear_reused(const char* prefix, int k) { + Module module(g_dir + "/" + prefix + ".pte"); + ASSERT_EQ(module.load_forward(), Error::Ok) << "load " << prefix << ".pte"; + for (int m_rows : {128, 32, 1, 128}) { + run_linear(module, m_rows, prefix, kFp32LinearN, k, 1e-3f); + } +} + +constexpr int kQkvNq = 2048; +constexpr int kQkvNk = 512; +constexpr int kQkvNv = 512; +constexpr int kQkvMaxM = 16; + +void run_qkv_routes(Module& m, int m_rows) { + const std::string base = + g_dir + "/qkv_routes.S" + std::to_string(m_rows) + "."; + auto input = read_bin(base + "input.bin"); + ASSERT_FALSE(input.empty()) << "missing qkv_routes.S" << m_rows; + auto tensor = make_tensor_ptr({m_rows, kLinK}, std::move(input)); + auto result = m.forward({EValue(tensor)}); + ASSERT_TRUE(result.ok()) << "qkv_routes M=" << m_rows << " forward failed"; + ASSERT_EQ(result.get().size(), 3u); + + const int widths[] = {kQkvNq, kQkvNk, kQkvNv}; + const char* names[] = {"q", "k", "v"}; + for (size_t i = 0; i < 3; i++) { + ASSERT_TRUE(result.get()[i].isTensor()); + const auto& output = result.get()[i].toTensor(); + ASSERT_EQ(output.dim(), 2); + ASSERT_EQ(output.size(0), m_rows); + ASSERT_EQ(output.size(1), widths[i]); + const size_t numel = static_cast(m_rows) * widths[i]; + std::vector got( + output.const_data_ptr(), output.const_data_ptr() + numel); + auto golden = read_bin(base + names[i] + ".bin"); + ASSERT_EQ(golden.size(), numel); + EXPECT_LT(max_err(got, golden), 1e-2f) + << "qkv_routes " << names[i] << " M=" << m_rows; + } +} + +constexpr int kQkvBk64K = 2048; +constexpr int kQkvBk64MaxM = 128; + +void run_qkv_bk64_routes(Module& module, int m_rows) { + const std::string path = + g_dir + "/qkv_bk64_routes.S" + std::to_string(m_rows) + ".input.bin"; + auto input = read_bin(path); + ASSERT_EQ(input.size(), static_cast(m_rows) * kQkvBk64K); + auto tensor = make_tensor_ptr({m_rows, kQkvBk64K}, std::move(input)); + auto result = module.forward({EValue(tensor)}); + ASSERT_TRUE(result.ok()) << "qkv_bk64_routes M=" << m_rows; + ASSERT_EQ(result.get().size(), 3u); + const int widths[] = {kQkvNq, kQkvNk, kQkvNv}; + for (size_t i = 0; i < 3; i++) { + ASSERT_TRUE(result.get()[i].isTensor()); + const auto& output = result.get()[i].toTensor(); + EXPECT_EQ( + static_cast(output.numel()), + static_cast(m_rows) * widths[i]); + } +} + +constexpr int kBk64K = 2048; +constexpr int kBk64N = 2048; +constexpr int kBk64KvN = 512; +constexpr int kBk64GateN = 8192; +constexpr int kBk64DownK = 8192; +constexpr float kBk64Atol = 5e-2f; +constexpr float kBk64Rtol = 3e-2f; +constexpr float kBk64Nrmse = 7.5e-3f; +constexpr float kBk64TailNrmse = 1e-2f; +constexpr float kBk64DownAtol = 5e-2f; +constexpr float kBk64DownRtol = 8e-2f; +constexpr float kBk64DownNrmse = 1.5e-2f; +constexpr float kBk64DownTailNrmse = 2e-2f; + +void expect_bk64_tensor( + const executorch::aten::Tensor& output, + const std::vector& golden, + int m_rows, + int width, + const std::string& label) { + const size_t numel = static_cast(m_rows) * width; + ASSERT_EQ(static_cast(output.numel()), numel) << label; + ASSERT_EQ(golden.size(), numel) << label; + const float* data = output.const_data_ptr(); + double error_sq_sum = 0.0; + double golden_sq_sum = 0.0; + double tail_error_sq_sum = 0.0; + double tail_golden_sq_sum = 0.0; + bool within_tolerance = true; + const size_t tail_begin = static_cast(m_rows - 1) * width; + for (size_t i = 0; i < numel; ++i) { + ASSERT_TRUE(std::isfinite(data[i])) << label << " i=" << i; + ASSERT_TRUE(std::isfinite(golden[i])) << label << " golden i=" << i; + const double error = static_cast(data[i]) - golden[i]; + const float abs_error = std::fabs(static_cast(error)); + const float rel_error = abs_error / std::fmax(std::fabs(golden[i]), 1e-6f); + within_tolerance &= abs_error <= kBk64Atol || rel_error <= kBk64Rtol; + error_sq_sum += error * error; + golden_sq_sum += static_cast(golden[i]) * golden[i]; + if (i >= tail_begin) { + tail_error_sq_sum += error * error; + tail_golden_sq_sum += static_cast(golden[i]) * golden[i]; + } + } + EXPECT_TRUE(within_tolerance) << label << " hybrid tolerance"; + ASSERT_GT(golden_sq_sum, 0.0) << label << " zero golden norm"; + EXPECT_LT(std::sqrt(error_sq_sum / golden_sq_sum), kBk64Nrmse) + << label << " full-output NRMSE"; + ASSERT_GT(tail_golden_sq_sum, 0.0) << label << " zero final-row golden norm"; + EXPECT_LT(std::sqrt(tail_error_sq_sum / tail_golden_sq_sum), kBk64TailNrmse) + << label << " final-row NRMSE"; + + for (size_t i : + {size_t{0}, static_cast(width - 1), tail_begin, numel - 1}) { + const float abs_error = std::fabs(data[i] - golden[i]); + const float rel_error = abs_error / std::fmax(std::fabs(golden[i]), 1e-6f); + EXPECT_TRUE(abs_error <= kBk64Atol || rel_error <= kBk64Rtol) + << label << " boundary i=" << i; + } +} + +void run_bk64_qkv( + Module& module, + int m_rows, + const char* prefix, + int q_width = kBk64N, + int k_width = kBk64KvN, + int v_width = kBk64KvN, + bool separate_v_input = false) { + const std::string base = + g_dir + "/" + prefix + ".S" + std::to_string(m_rows) + "."; + auto input = read_bin(base + "input.bin"); + ASSERT_EQ(input.size(), static_cast(m_rows) * kBk64K); + auto input_tensor = make_tensor_ptr({m_rows, kBk64K}, std::move(input)); + std::vector inputs{EValue(input_tensor)}; + decltype(input_tensor) v_input_tensor; + if (separate_v_input) { + auto v_input = read_bin(base + "v_input.bin"); + ASSERT_EQ(v_input.size(), static_cast(m_rows) * kBk64K); + v_input_tensor = make_tensor_ptr({m_rows, kBk64K}, std::move(v_input)); + inputs.emplace_back(v_input_tensor); + } + auto result = module.forward(inputs); + ASSERT_TRUE(result.ok()) << prefix << " M=" << m_rows; + ASSERT_EQ(result.get().size(), 3) << prefix << " M=" << m_rows; + const int widths[] = {q_width, k_width, v_width}; + const char* names[] = {"q", "k", "v"}; + for (size_t i = 0; i < 3; ++i) { + ASSERT_TRUE(result.get()[i].isTensor()) << prefix << " output " << names[i]; + expect_bk64_tensor( + result.get()[i].toTensor(), + read_bin(base + names[i] + ".bin"), + m_rows, + widths[i], + std::string(prefix) + " M=" + std::to_string(m_rows) + " " + names[i]); + } +} + +void run_bk64_linear( + Module& module, + int m_rows, + const char* prefix, + int k = kBk64K, + int n = kBk64N, + float atol = kBk64Atol, + float rtol = kBk64Rtol, + float nrmse = kBk64Nrmse, + float tail_nrmse = kBk64TailNrmse) { + run_linear(module, m_rows, prefix, n, k, atol, rtol, nrmse, tail_nrmse); +} + +constexpr int kSwiGluWidth = 8192; +constexpr int kSwiGluSmallWidth = 64; +constexpr int kSwiGluQkvOverlapWidth = 512; +constexpr int kSwiGluK = 64; + +void run_swiglu( + Module& module, + int m_rows, + const char* prefix, + int width, + bool separate_inputs = false) { + const std::string base = + g_dir + "/" + prefix + ".S" + std::to_string(m_rows) + "."; + auto input = read_bin(base + "input.bin"); + auto golden = read_bin(base + "golden.bin"); + ASSERT_FALSE(input.empty() || golden.empty()) + << "missing " << prefix << ".S" << m_rows; + auto input_tensor = make_tensor_ptr({m_rows, kSwiGluK}, std::move(input)); + std::vector inputs{EValue(input_tensor)}; + decltype(input_tensor) up_input_tensor; + if (separate_inputs) { + auto up_input = read_bin(base + "up_input.bin"); + ASSERT_FALSE(up_input.empty()); + up_input_tensor = make_tensor_ptr({m_rows, kSwiGluK}, std::move(up_input)); + inputs.emplace_back(up_input_tensor); + } + auto result = module.forward(inputs); + ASSERT_TRUE( + result.ok() && result.get().size() == 1 && result.get()[0].isTensor()) + << prefix << " M=" << m_rows << " forward failed"; + const auto& output = result.get()[0].toTensor(); + const size_t numel = static_cast(m_rows) * width; + ASSERT_EQ(static_cast(output.numel()), numel); + std::vector got( + output.const_data_ptr(), output.const_data_ptr() + numel); + EXPECT_LT(max_err(got, golden), 1e-2f) << prefix << " M=" << m_rows; +} + +void run_swiglu_outputs( + Module& module, + int m_rows, + const char* prefix, + size_t output_count) { + const std::string base = + g_dir + "/" + prefix + ".S" + std::to_string(m_rows) + "."; + auto input = read_bin(base + "input.bin"); + ASSERT_FALSE(input.empty()); + auto input_tensor = make_tensor_ptr({m_rows, kSwiGluK}, std::move(input)); + auto result = module.forward({EValue(input_tensor)}); + ASSERT_TRUE(result.ok()); + ASSERT_EQ(result.get().size(), output_count); + const size_t numel = static_cast(m_rows) * kSwiGluSmallWidth; + for (size_t index = 0; index < result.get().size(); ++index) { + ASSERT_TRUE(result.get()[index].isTensor()); + const auto& output = result.get()[index].toTensor(); + ASSERT_EQ(static_cast(output.numel()), numel); + std::vector got( + output.const_data_ptr(), output.const_data_ptr() + numel); + const auto golden = + read_bin(base + "golden" + std::to_string(index) + ".bin"); + EXPECT_LT(max_err(got, golden), 1e-2f) << "output " << index; + } +} + +void run_swiglu_graph_outputs(Module& module, int m_rows) { + run_swiglu_outputs(module, m_rows, "dyn_swiglu_graph_outputs", 4); +} + +void run_swiglu_qkv_overlap(Module& module, int m_rows) { + const std::string base = + g_dir + "/dyn_swiglu_qkv_overlap.S" + std::to_string(m_rows) + "."; + auto input = read_bin(base + "input.bin"); + ASSERT_FALSE(input.empty()); + auto input_tensor = make_tensor_ptr({m_rows, kBk64K}, std::move(input)); + auto result = module.forward({EValue(input_tensor)}); + ASSERT_TRUE(result.ok()); + ASSERT_EQ(result.get().size(), 2u); + const int widths[] = {2048, kSwiGluQkvOverlapWidth}; + for (size_t index = 0; index < 2; index++) { + ASSERT_TRUE(result.get()[index].isTensor()); + const auto& output = result.get()[index].toTensor(); + const size_t numel = static_cast(m_rows) * widths[index]; + ASSERT_EQ(static_cast(output.numel()), numel); + std::vector got( + output.const_data_ptr(), output.const_data_ptr() + numel); + const auto golden = + read_bin(base + "golden" + std::to_string(index) + ".bin"); + EXPECT_LT(max_err(got, golden), 1e-2f) << "overlap output " << index; + } } // Dynamic SDPA (GQA prefill, input_pos=0): q[1,s,hq,d] k/v[1,s,hkv,d] // caches[1,cmax,hkv,d]; attn output [1,s,hq,d] selected by shape (3 outputs). constexpr int kSdHq = 8, kSdHkv = 2, kSdD = 16, kSdCmax = 64; -void check_sdpa(int s) { - Module m(g_dir + "/sdpa_dyn.pte"); - ASSERT_EQ(m.load_forward(), Error::Ok) << "sdpa_dyn S=" << s << " load"; - const std::string b = g_dir + "/sdpa_dyn.S" + std::to_string(s) + "."; +void run_sdpa_case( + Module& m, + int s, + const char* prefix, + int hq, + int hkv, + int d, + int cmax, + float max_error_limit = 2e-3f) { + const std::string b = g_dir + "/" + prefix + ".S" + std::to_string(s) + "."; auto q = read_bin(b + "q.bin"); auto k = read_bin(b + "k.bin"); auto v = read_bin(b + "v.bin"); @@ -158,34 +522,230 @@ void check_sdpa(int s) { q.empty() || k.empty() || v.empty() || kc.empty() || vc.empty() || golden.empty()) << "missing sdpa_dyn.S" << s; - auto tq = make_tensor_ptr({1, s, kSdHq, kSdD}, std::move(q)); - auto tk = make_tensor_ptr({1, s, kSdHkv, kSdD}, std::move(k)); - auto tv = make_tensor_ptr({1, s, kSdHkv, kSdD}, std::move(v)); - auto tkc = make_tensor_ptr({1, kSdCmax, kSdHkv, kSdD}, std::move(kc)); - auto tvc = make_tensor_ptr({1, kSdCmax, kSdHkv, kSdD}, std::move(vc)); + auto tq = make_tensor_ptr({1, s, hq, d}, std::move(q)); + auto tk = make_tensor_ptr({1, s, hkv, d}, std::move(k)); + auto tv = make_tensor_ptr({1, s, hkv, d}, std::move(v)); + auto tkc = make_tensor_ptr({1, cmax, hkv, d}, std::move(kc)); + auto tvc = make_tensor_ptr({1, cmax, hkv, d}, std::move(vc)); auto r = m.forward({EValue(tq), EValue(tk), EValue(tv), EValue(tkc), EValue(tvc)}); ASSERT_TRUE(r.ok()) << "sdpa S=" << s << " forward failed (err=" << (int)r.error() << ")"; // Select the attn output by full shape [1,s,hq,d] (never numel). const float* attn = nullptr; - const size_t numel = static_cast(s) * kSdHq * kSdD; + const size_t numel = static_cast(s) * hq * d; for (size_t i = 0; i < r.get().size(); i++) { if (!r.get()[i].isTensor()) { continue; } const auto& t = r.get()[i].toTensor(); - if (t.dim() == 4 && t.size(1) == s && t.size(2) == kSdHq && - t.size(3) == kSdD) { + if (t.dim() == 4 && t.size(1) == s && t.size(2) == hq && t.size(3) == d) { attn = t.const_data_ptr(); break; } } ASSERT_NE(attn, nullptr) << "sdpa S=" << s << ": no attn output of shape [1," - << s << "," << kSdHq << "," << kSdD << "]"; + << s << "," << hq << "," << d << "]"; std::vector got(attn, attn + numel); const float e = max_err(got, golden); - EXPECT_LT(e, 2e-3f) << "sdpa_dyn S=" << s << " max_err=" << e; + EXPECT_LT(e, max_error_limit) + << prefix << " S=" << s << " full-output max_err=" << e; +} + +void run_sdpa(Module& m, int s) { + run_sdpa_case(m, s, "sdpa_dyn", kSdHq, kSdHkv, kSdD, kSdCmax); +} + +void check_sdpa(int s) { + Module m(g_dir + "/sdpa_dyn.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "sdpa_dyn S=" << s << " load"; + run_sdpa(m, s); +} + +constexpr int kK16Hq = 32; +constexpr int kK16Hkv = 8; +constexpr int kK16D = 64; +constexpr int kQwen3Hq = 16; +constexpr int kQwen3Hkv = 8; +constexpr int kQwen3D = 128; + +bool k16_device_supported() { + const auto* context = get_default_webgpu_context(); + WGPULimits limits = {}; + return context != nullptr && context->shader_f16_supported && + wgpuDeviceGetLimits(context->device, &limits) == WGPUStatus_Success && + limits.maxComputeInvocationsPerWorkgroup >= 128u && + limits.maxComputeWorkgroupSizeX >= 32u && + limits.maxComputeWorkgroupSizeY >= 4u && + limits.maxComputeWorkgroupStorageSize >= 14720u; +} + +bool qwen3_q16_device_supported() { + constexpr uint32_t kQ16StorageBytes = 512u * 4u * sizeof(float) + + 512u * 4u * sizeof(uint16_t) + 128u * 2u * sizeof(float) + + 3u * 16u * sizeof(float); + const auto* context = get_default_webgpu_context(); + WGPULimits limits = {}; + return context != nullptr && context->shader_f16_supported && + wgpuDeviceGetLimits(context->device, &limits) == WGPUStatus_Success && + limits.maxComputeWorkgroupSizeX >= 16u && + limits.maxComputeWorkgroupSizeY >= 8u && + limits.maxComputeInvocationsPerWorkgroup >= 128u && + limits.maxComputeWorkgroupStorageSize >= kQ16StorageBytes && + limits.maxStorageBuffersPerShaderStage >= 4u; +} + +void load_sdpa_module(Module& module, bool f16_kv) { + if (!f16_kv) { + ASSERT_EQ(module.load_forward(), Error::Ok); + return; + } + executorch::runtime::BackendOptions<1> options; + ASSERT_EQ(options.set_option("enable_f16_kv_cache", true), Error::Ok); + executorch::runtime::LoadBackendOptionsMap option_map; + ASSERT_EQ(option_map.set_options("VulkanBackend", options.view()), Error::Ok); + ASSERT_EQ(module.load_forward(nullptr, nullptr, &option_map), Error::Ok); +} + +void run_k16_sdpa( + Module& module, + int s, + const char* prefix, + int hq = kK16Hq, + int hkv = kK16Hkv, + int d = kK16D, + bool prime = false, + float max_error_limit = 3e-3f, + bool initial = false) { + const std::string suffix = initial ? ".initial." + : prime ? ".prime." + : ".S" + std::to_string(s) + "."; + const std::string base = g_dir + "/" + prefix + suffix; + auto q = read_bin(base + "q.bin"); + auto k = read_bin(base + "k.bin"); + auto v = read_bin(base + "v.bin"); + auto control = read_bin(base + "control.bin"); + auto golden = read_bin(base + "golden.bin"); + ASSERT_FALSE( + q.empty() || k.empty() || v.empty() || control.empty() || golden.empty()); + auto tq = make_tensor_ptr({1, s, hq, d}, std::move(q)); + auto tk = make_tensor_ptr({1, s, hkv, d}, std::move(k)); + auto tv = make_tensor_ptr({1, s, hkv, d}, std::move(v)); + const int control_size = static_cast(control.size()); + auto tcontrol = make_tensor_ptr({1, control_size}, std::move(control)); + auto result = + module.forward({EValue(tq), EValue(tk), EValue(tv), EValue(tcontrol)}); + ASSERT_TRUE( + result.ok() && result.get().size() == 1 && result.get()[0].isTensor()) + << prefix << " S=" << s << " forward failed"; + const auto& output = result.get()[0].toTensor(); + const size_t token_width = static_cast(hq) * d; + const size_t numel = static_cast(s) * token_width; + ASSERT_EQ(static_cast(output.numel()), numel); + std::vector got( + output.const_data_ptr(), output.const_data_ptr() + numel); + ASSERT_EQ(got.size(), golden.size()); + EXPECT_LT(max_err(got, golden), max_error_limit) + << prefix << " S=" << s << " full output"; + const std::set tokens = { + 0, std::min(15, s - 1), std::min(16, s - 1), s - 1}; + for (int token : tokens) { + float token_error = 0.0f; + const size_t begin = static_cast(token) * token_width; + for (size_t i = begin; i < begin + token_width; ++i) { + token_error = std::fmax(token_error, std::fabs(got[i] - golden[i])); + } + EXPECT_LT(token_error, max_error_limit) + << prefix << " S=" << s << " causal token=" << token; + } +} + +void prime_k16_sdpa( + Module& module, + const char* prefix, + int hq = kK16Hq, + int hkv = kK16Hkv, + int d = kK16D, + float max_error_limit = 3e-3f) { + run_k16_sdpa(module, 12, prefix, hq, hkv, d, true, max_error_limit); +} + +#ifdef WGPU_BACKEND_ENABLE_PROFILING +void expect_sdpa_route( + const std::vector& names, + int s, + bool expect_k16, + const char* k16_kernel_name = "sdpa_streaming_attention_k16_causal_bound") { + const bool expect_fd = s == 1; + const bool expect_materialized = !expect_fd && !expect_k16; + EXPECT_EQ(std::count(names.begin(), names.end(), "update_cache"), 2); + EXPECT_EQ( + std::count(names.begin(), names.end(), k16_kernel_name), + expect_k16 ? 1 : 0); + EXPECT_EQ( + std::count(names.begin(), names.end(), "fd_split"), expect_fd ? 1 : 0); + EXPECT_EQ( + std::count(names.begin(), names.end(), "fd_reduce"), expect_fd ? 1 : 0); + EXPECT_EQ( + std::count(names.begin(), names.end(), "sdpa_compute_attn_weights"), + expect_materialized ? 1 : 0); + EXPECT_EQ( + std::count(names.begin(), names.end(), "sdpa_softmax"), + expect_materialized ? 1 : 0); + EXPECT_EQ( + std::count(names.begin(), names.end(), "sdpa_compute_out"), + expect_materialized ? 1 : 0); + EXPECT_EQ( + names.size(), + static_cast(2 + (expect_k16 ? 1 : (expect_fd ? 2 : 3)))); +} +#endif + +void run_combined_routes(Module& m, int s) { + const std::string b = g_dir + "/combined_routes.S" + std::to_string(s) + "."; + auto x = read_bin(b + "x.bin"); + auto q = read_bin(b + "q.bin"); + auto k = read_bin(b + "k.bin"); + auto v = read_bin(b + "v.bin"); + auto kc = read_bin(b + "kc.bin"); + auto vc = read_bin(b + "vc.bin"); + auto golden = read_bin(b + "golden.bin"); + ASSERT_FALSE( + x.empty() || q.empty() || k.empty() || v.empty() || kc.empty() || + vc.empty() || golden.empty()) + << "missing combined_routes.S" << s; + auto tx = make_tensor_ptr({s, kLinK}, std::move(x)); + auto tq = make_tensor_ptr({1, s, kSdHq, kSdD}, std::move(q)); + auto tk = make_tensor_ptr({1, s, kSdHkv, kSdD}, std::move(k)); + auto tv = make_tensor_ptr({1, s, kSdHkv, kSdD}, std::move(v)); + auto tkc = make_tensor_ptr({1, kSdCmax, kSdHkv, kSdD}, std::move(kc)); + auto tvc = make_tensor_ptr({1, kSdCmax, kSdHkv, kSdD}, std::move(vc)); + auto result = m.forward( + {EValue(tx), + EValue(tq), + EValue(tk), + EValue(tv), + EValue(tkc), + EValue(tvc)}); + ASSERT_TRUE(result.ok()) << "combined routes S=" << s + << " forward failed (err=" << (int)result.error() + << ")"; + const float* attn = nullptr; + const size_t numel = static_cast(s) * kSdHq * kSdD; + for (const auto& output : result.get()) { + if (!output.isTensor()) { + continue; + } + const auto& tensor = output.toTensor(); + if (tensor.dim() == 4 && tensor.size(1) == s && tensor.size(2) == kSdHq && + tensor.size(3) == kSdD) { + attn = tensor.const_data_ptr(); + break; + } + } + ASSERT_NE(attn, nullptr); + const std::vector got(attn, attn + numel); + EXPECT_LT(max_err(got, golden), 1e-2f) << "combined_routes S=" << s; } // Dynamic embedding: int64 token ids [N] -> [N, kEmbDim] fp32. The int64 host @@ -390,6 +950,23 @@ TEST(DynamicShape, RmsMul) { } } +// I0: dynamic fp32 linear preserves bias across repeated resizes. +TEST(DynamicShape, Fp32LinearVec4BiasedReusedGraph) { + check_fp32_linear_reused("dyn_linear_fp32_vec4_bias", 64); +} + +TEST(DynamicShape, Fp32LinearVec4UnbiasedReusedGraph) { + check_fp32_linear_reused("dyn_linear_fp32_vec4_no_bias", 64); +} + +TEST(DynamicShape, Fp32LinearTiledBiasedReusedGraph) { + check_fp32_linear_reused("dyn_linear_fp32_tiled_bias", 63); +} + +TEST(DynamicShape, Fp32LinearTiledUnbiasedReusedGraph) { + check_fp32_linear_reused("dyn_linear_fp32_tiled_no_bias", 63); +} + // I: dynamic 4-bit quantized linear (prefill GEMM) at several M. TEST(DynamicShape, QuantizedLinear) { for (int m_rows : {128, 32, 1}) { @@ -407,9 +984,7 @@ TEST(DynamicShape, QuantizedLinearReusedGraph) { } } -// I3: dynamic linear at N=2048 -> the shmem-GEMM route (K>=4096||N>=2048); the -// resize hook recomputes the shmem tile count for the live M on the fixed shmem -// pipeline (M=1 exercises a partial row-tile). +// I3: K=72 disables Steel; N=2048 forces shmem for M>1 and bicol for M=1. TEST(DynamicShape, QuantizedLinearShmem) { for (int m_rows : {128, 32, 1}) { check_linear_shmem(m_rows); @@ -421,26 +996,688 @@ TEST(DynamicShape, QuantizedLinearShmemReusedGraph) { Module m(g_dir + "/dyn_linear_shmem.pte"); ASSERT_EQ(m.load_forward(), Error::Ok) << "load dyn_linear_shmem.pte"; for (int m_rows : {128, 32, 1, 128}) { - run_linear(m, m_rows, "dyn_linear_shmem", kLinNShmem); + run_linear(m, m_rows, "dyn_linear_shmem", kLinNShmem, kLinAltK); } } -// J: dynamic SDPA (GQA prefill) at several seq-len S. The whole case skips -// while op coverage is pending (the dynamic-S build throws err 48 until -// registered). -TEST(DynamicShape, Sdpa) { - { - Module probe(g_dir + "/sdpa_dyn.pte"); - if (probe.load_forward() == Error::DelegateInvalidCompatibility) { - GTEST_SKIP() << "sdpa_dyn pending op coverage (err " - << (int)Error::DelegateInvalidCompatibility << ")"; +// I5: K=72 disables Steel; N=128 keeps the tiled fallback for M>1. +TEST(DynamicShape, QuantizedLinearTiled) { + for (int m_rows : {128, 32, 1}) { + check_linear_tiled(m_rows); + } +} + +TEST(DynamicShape, QuantizedLinearTiledReusedGraph) { + Module m(g_dir + "/dyn_linear_tiled.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "load dyn_linear_tiled.pte"; + for (int m_rows : {128, 1, 32, 1, 128}) { + run_linear(m, m_rows, "dyn_linear_tiled", kLinN, kLinAltK); + } +} + +// The first max-shape execution does not invoke resize hooks. This sequence +// therefore checks both initial QKV route selection and max/decode transitions. +TEST(DynamicShape, QkvRoutesReusedGraph) { + Module m(g_dir + "/qkv_routes.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "load qkv_routes.pte"; + for (int m_rows : {kQkvMaxM, 1, kQkvMaxM}) { + run_qkv_routes(m, m_rows); + } +} + +TEST(DynamicShape, QuantizedLinearBk64ReusedGraphAndFallbacks) { + Module candidate(g_dir + "/dyn_linear_bk64.pte"); + ASSERT_EQ(candidate.load_forward(), Error::Ok) << "load dyn_linear_bk64.pte"; + for (int m_rows : {512, 511, 508, 128, 127, 1, 508, 512}) { + run_bk64_linear(candidate, m_rows, "dyn_linear_bk64"); + } + + Module gate(g_dir + "/dyn_linear_bk64_gate.pte"); + ASSERT_EQ(gate.load_forward(), Error::Ok); + for (int m_rows : {512, 508, 128}) { + run_bk64_linear(gate, m_rows, "dyn_linear_bk64_gate", kBk64K, kBk64GateN); + } + + Module down(g_dir + "/dyn_linear_bk64_down.pte"); + ASSERT_EQ(down.load_forward(), Error::Ok); + for (int m_rows : {512, 508, 128}) { + run_bk64_linear( + down, + m_rows, + "dyn_linear_bk64_down", + kBk64DownK, + kBk64N, + kBk64DownAtol, + kBk64DownRtol, + kBk64DownNrmse, + kBk64DownTailNrmse); + } + + Module group32(g_dir + "/dyn_linear_bk64_group32.pte"); + ASSERT_EQ(group32.load_forward(), Error::Ok); + run_bk64_linear(group32, 128, "dyn_linear_bk64_group32"); + + Module bias(g_dir + "/dyn_linear_bk64_bias.pte"); + ASSERT_EQ(bias.load_forward(), Error::Ok); + run_bk64_linear(bias, 128, "dyn_linear_bk64_bias"); + + Module kv_shape(g_dir + "/dyn_linear_bk64_kv_shape.pte"); + ASSERT_EQ(kv_shape.load_forward(), Error::Ok); + run_bk64_linear(kv_shape, 128, "dyn_linear_bk64_kv_shape", kBk64K, kBk64KvN); +} + +TEST(DynamicShape, QuantizedLinearBk64QkvReusedGraphAndFallbacks) { + Module candidate(g_dir + "/dyn_qkv_bk64.pte"); + ASSERT_EQ(candidate.load_forward(), Error::Ok) << "load dyn_qkv_bk64.pte"; + for (int m_rows : {512, 511, 508, 128, 127, 16, 2, 1, 508, 512}) { + run_bk64_qkv(candidate, m_rows, "dyn_qkv_bk64"); + } + + Module group32(g_dir + "/dyn_qkv_bk64_group32.pte"); + ASSERT_EQ(group32.load_forward(), Error::Ok); + run_bk64_qkv(group32, 128, "dyn_qkv_bk64_group32"); + + Module bias(g_dir + "/dyn_qkv_bk64_bias.pte"); + ASSERT_EQ(bias.load_forward(), Error::Ok); + run_bk64_qkv(bias, 128, "dyn_qkv_bk64_bias"); + + Module wrong_width(g_dir + "/dyn_qkv_bk64_wrong_width.pte"); + ASSERT_EQ(wrong_width.load_forward(), Error::Ok); + run_bk64_qkv( + wrong_width, 128, "dyn_qkv_bk64_wrong_width", kBk64N, kBk64N, kBk64KvN); + + Module different_input(g_dir + "/dyn_qkv_bk64_different_input.pte"); + ASSERT_EQ(different_input.load_forward(), Error::Ok); + run_bk64_qkv( + different_input, + 128, + "dyn_qkv_bk64_different_input", + kBk64N, + kBk64KvN, + kBk64KvN, + true); +} + +#ifdef WGPU_BACKEND_ENABLE_PROFILING +TEST(DynamicShape, QkvLiveRoutesProfile) { + const auto* context = get_default_webgpu_context(); + if (std::getenv("WEBGPU_TIMESTAMP_QUERY") == nullptr || context == nullptr || + !context->timestamp_supported || !context->shader_f16_supported) { + GTEST_SKIP() << "timestamp queries or shader-f16 unavailable"; + } + Module m(g_dir + "/qkv_routes.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "load qkv_routes.pte"; + for (int m_rows : {kQkvMaxM, 1, kQkvMaxM}) { + run_qkv_routes(m, m_rows); + const auto names = current_profile_names(); + EXPECT_EQ( + std::count(names.begin(), names.end(), "linear_q4gsw_bk64_qkv"), 0); + EXPECT_EQ( + std::count(names.begin(), names.end(), "linear_q4gsw_coop4_bicol"), + m_rows == 1 ? 3 : 0); + EXPECT_EQ( + std::count_if( + names.begin(), + names.end(), + [](const std::string& name) { + return name.rfind("linear_q4gsw", 0) == 0; + }), + 3); + } +} + +TEST(DynamicShape, QkvBk64LiveRoutesProfile) { + const auto* context = get_default_webgpu_context(); + if (std::getenv("WEBGPU_TIMESTAMP_QUERY") == nullptr || context == nullptr || + !context->timestamp_supported || !context->shader_f16_supported) { + GTEST_SKIP() << "timestamp queries or shader-f16 unavailable"; + } + WGPULimits limits = {}; + if (wgpuDeviceGetLimits(context->device, &limits) != WGPUStatus_Success || + limits.maxComputeInvocationsPerWorkgroup < 256u || + limits.maxComputeWorkgroupSizeX < 16u || + limits.maxComputeWorkgroupSizeY < 16u || + limits.maxComputeWorkgroupStorageSize < 16384u || + limits.maxComputeWorkgroupsPerDimension < 384u) { + GTEST_SKIP() << "BK64 workgroup limits unavailable"; + } + Module module(g_dir + "/qkv_bk64_routes.pte"); + ASSERT_EQ(module.load_forward(), Error::Ok); + for (int m_rows : {kQkvBk64MaxM, 1, kQkvBk64MaxM}) { + run_qkv_bk64_routes(module, m_rows); + const auto names = current_profile_names(); + EXPECT_EQ( + std::count(names.begin(), names.end(), "linear_q4gsw_bk64_qkv"), + m_rows > 1 ? 1 : 0); + EXPECT_EQ( + std::count(names.begin(), names.end(), "linear_q4gsw_coop4_bicol"), + m_rows == 1 ? 3 : 0); + EXPECT_EQ(std::count(names.begin(), names.end(), "linear_q4gsw_bk64"), 0); + EXPECT_EQ(std::count(names.begin(), names.end(), "linear_q4gsw_steel"), 0); + EXPECT_EQ(std::count(names.begin(), names.end(), "linear_q4gsw_shmem"), 0); + EXPECT_EQ(std::count(names.begin(), names.end(), "linear_q4gsw_tiled"), 0); + } +} + +void expect_single_q4_profile( + const char* expected_name, + uint32_t expected_x, + const char* fixture, + int m_rows) { + const auto* context = get_default_webgpu_context(); + ASSERT_NE(context, nullptr); + ASSERT_NE(context->querypool, nullptr); + const auto& profile = context->querypool->results(); + const auto is_q4 = [](const auto& duration) { + return duration.kernel_name.rfind("linear_q4gsw", 0) == 0; + }; + ASSERT_EQ(std::count_if(profile.begin(), profile.end(), is_q4), 1) + << fixture << " M=" << m_rows; + const auto active = std::find_if(profile.begin(), profile.end(), is_q4); + ASSERT_NE(active, profile.end()); + EXPECT_EQ(active->kernel_name, expected_name) << fixture << " M=" << m_rows; + EXPECT_EQ(active->global_wg[0], expected_x) << fixture << " M=" << m_rows; + EXPECT_EQ(active->global_wg[1], 1u) << fixture << " M=" << m_rows; +} + +TEST(DynamicShape, QuantizedLinearBk64ProfileSoleWriter) { + const auto* context = get_default_webgpu_context(); + if (std::getenv("WEBGPU_TIMESTAMP_QUERY") == nullptr || context == nullptr || + !context->timestamp_supported || !context->shader_f16_supported) { + GTEST_SKIP() << "BK64 timestamp or shader-f16 capability unavailable"; + } + WGPULimits limits = {}; + if (wgpuDeviceGetLimits(context->device, &limits) != WGPUStatus_Success || + limits.maxComputeInvocationsPerWorkgroup < 256u || + limits.maxComputeWorkgroupStorageSize < 16384u || + limits.maxComputeWorkgroupsPerDimension < 1024u) { + GTEST_SKIP() << "BK64 workgroup limits unavailable"; + } + + Module candidate(g_dir + "/dyn_linear_bk64.pte"); + ASSERT_EQ(candidate.load_forward(), Error::Ok); + for (int m_rows : {512, 511, 508, 128, 127, 1, 508, 512}) { + run_bk64_linear(candidate, m_rows, "dyn_linear_bk64"); + const char* expected = m_rows == 1 ? "linear_q4gsw_coop4_bicol" + : (m_rows == 128 || m_rows == 508 || m_rows == 512) + ? "linear_q4gsw_bk64" + : "linear_q4gsw_steel"; + const uint32_t expected_x = m_rows == 1 ? 1024u + : (m_rows == 128 || m_rows == 127) ? 64u + : 256u; + expect_single_q4_profile(expected, expected_x, "dyn_linear_bk64", m_rows); + } + + Module gate(g_dir + "/dyn_linear_bk64_gate.pte"); + ASSERT_EQ(gate.load_forward(), Error::Ok); + for (int m_rows : {512, 508, 128}) { + run_bk64_linear(gate, m_rows, "dyn_linear_bk64_gate", kBk64K, kBk64GateN); + expect_single_q4_profile( + "linear_q4gsw_bk64", + m_rows == 128 ? 256u : 1024u, + "dyn_linear_bk64_gate", + m_rows); + } + + Module down(g_dir + "/dyn_linear_bk64_down.pte"); + ASSERT_EQ(down.load_forward(), Error::Ok); + for (int m_rows : {512, 508, 128}) { + run_bk64_linear( + down, + m_rows, + "dyn_linear_bk64_down", + kBk64DownK, + kBk64N, + kBk64DownAtol, + kBk64DownRtol, + kBk64DownNrmse, + kBk64DownTailNrmse); + expect_single_q4_profile( + "linear_q4gsw_bk64", + m_rows == 128 ? 64u : 256u, + "dyn_linear_bk64_down", + m_rows); + } + + struct NegativeRoute { + const char* fixture; + int n; + uint32_t expected_x; + }; + for (const auto& negative : std::vector{ + {"dyn_linear_bk64_group32", kBk64N, 64u}, + {"dyn_linear_bk64_bias", kBk64N, 64u}, + {"dyn_linear_bk64_kv_shape", kBk64KvN, 16u}}) { + Module module(g_dir + "/" + negative.fixture + ".pte"); + ASSERT_EQ(module.load_forward(), Error::Ok) << negative.fixture; + run_bk64_linear(module, 128, negative.fixture, kBk64K, negative.n); + expect_single_q4_profile( + "linear_q4gsw_steel", negative.expected_x, negative.fixture, 128); + } +} + +TEST(DynamicShape, QuantizedLinearBk64QkvProfileSoleWriter) { + const auto* context = get_default_webgpu_context(); + if (std::getenv("WEBGPU_TIMESTAMP_QUERY") == nullptr || context == nullptr || + !context->timestamp_supported || !context->shader_f16_supported) { + GTEST_SKIP() << "QKV timestamp or shader-f16 capability unavailable"; + } + WGPULimits limits = {}; + if (wgpuDeviceGetLimits(context->device, &limits) != WGPUStatus_Success || + limits.maxComputeInvocationsPerWorkgroup < 256u || + limits.maxComputeWorkgroupSizeX < 16u || + limits.maxComputeWorkgroupSizeY < 16u || + limits.maxComputeWorkgroupStorageSize < 16384u || + limits.maxComputeWorkgroupsPerDimension < 384u) { + GTEST_SKIP() << "QKV workgroup limits unavailable"; + } + const auto q4_profiles = [&]() { + std::vector names; + for (const auto& duration : context->querypool->results()) { + if (duration.kernel_name.rfind("linear_q4gsw", 0) == 0) { + names.push_back(duration.kernel_name); + } } + return names; + }; + + Module candidate(g_dir + "/dyn_qkv_bk64.pte"); + ASSERT_EQ(candidate.load_forward(), Error::Ok); + for (int m_rows : {512, 508, 128}) { + run_bk64_qkv(candidate, m_rows, "dyn_qkv_bk64"); + const auto names = q4_profiles(); + ASSERT_EQ(names.size(), 1) << "M=" << m_rows; + EXPECT_EQ(names[0], "linear_q4gsw_bk64_qkv") << "M=" << m_rows; + const auto& profile = context->querypool->results(); + const auto active = + std::find_if(profile.begin(), profile.end(), [](const auto& d) { + return d.kernel_name == "linear_q4gsw_bk64_qkv"; + }); + ASSERT_NE(active, profile.end()); + EXPECT_EQ(active->global_wg[0], m_rows == 128 ? 96u : 384u); + EXPECT_EQ(active->global_wg[1], 1u); + } + + for (int m_rows : {511, 127, 16, 2}) { + run_bk64_qkv(candidate, m_rows, "dyn_qkv_bk64"); + const auto names = q4_profiles(); + ASSERT_EQ(names.size(), 3) << "M=" << m_rows; + EXPECT_FALSE(contains_name(names, "linear_q4gsw_bk64_qkv")); + EXPECT_EQ(std::count(names.begin(), names.end(), "linear_q4gsw_steel"), 3) + << "M=" << m_rows; + } + + run_bk64_qkv(candidate, 1, "dyn_qkv_bk64"); + const auto decode_names = q4_profiles(); + ASSERT_EQ(decode_names.size(), 3); + EXPECT_EQ( + std::count( + decode_names.begin(), decode_names.end(), "linear_q4gsw_coop4_bicol"), + 3); + EXPECT_FALSE(contains_name(decode_names, "linear_q4gsw_bk64_qkv")); + + for (int m_rows : {508, 512}) { + run_bk64_qkv(candidate, m_rows, "dyn_qkv_bk64"); + const auto names = q4_profiles(); + ASSERT_EQ(names.size(), 1) << "re-entry M=" << m_rows; + EXPECT_EQ(names[0], "linear_q4gsw_bk64_qkv") << "re-entry M=" << m_rows; } + + struct NegativeRoute { + const char* fixture; + int q_width; + int k_width; + bool separate_v_input; + }; + for (const auto& negative : std::vector{ + {"dyn_qkv_bk64_group32", kBk64N, kBk64KvN, false}, + {"dyn_qkv_bk64_bias", kBk64N, kBk64KvN, false}, + {"dyn_qkv_bk64_wrong_width", kBk64N, kBk64N, false}, + {"dyn_qkv_bk64_different_input", kBk64N, kBk64KvN, true}}) { + Module module(g_dir + "/" + negative.fixture + ".pte"); + ASSERT_EQ(module.load_forward(), Error::Ok) << negative.fixture; + run_bk64_qkv( + module, + 128, + negative.fixture, + negative.q_width, + negative.k_width, + kBk64KvN, + negative.separate_v_input); + const auto names = q4_profiles(); + ASSERT_EQ(names.size(), 3) << negative.fixture; + EXPECT_FALSE(contains_name(names, "linear_q4gsw_bk64_qkv")) + << negative.fixture; + } +} + +TEST(DynamicShape, CombinedLiveRoutesProfile) { + const auto* context = get_default_webgpu_context(); + if (std::getenv("WEBGPU_TIMESTAMP_QUERY") == nullptr || context == nullptr || + !context->timestamp_supported || !k16_device_supported()) { + GTEST_SKIP() << "timestamp queries or K16 device limits unavailable"; + } + Module m(g_dir + "/combined_routes.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "load combined_routes.pte"; + for (int s : {64, 1, 16, 1, 64}) { + run_combined_routes(m, s); + const auto names = current_profile_names(); + ASSERT_EQ(names.size(), s == 1 ? 6 : 7); + EXPECT_EQ(std::count(names.begin(), names.end(), "add"), 1); + EXPECT_EQ( + std::count(names.begin(), names.end(), "linear_q4gsw_coop4_bicol"), + s == 1 ? 1 : 0); + EXPECT_EQ( + std::count(names.begin(), names.end(), "linear_q4gsw_steel"), + s != 1 ? 1 : 0); + EXPECT_FALSE(contains_name(names, "linear_q4gsw_shmem")); + EXPECT_FALSE(contains_name(names, "linear_q4gsw_tiled")); + EXPECT_EQ(std::count(names.begin(), names.end(), "update_cache"), 2); + EXPECT_EQ(std::count(names.begin(), names.end(), "fd_split"), s == 1); + EXPECT_EQ(std::count(names.begin(), names.end(), "fd_reduce"), s == 1); + EXPECT_EQ( + std::count(names.begin(), names.end(), "sdpa_compute_attn_weights"), + s != 1); + EXPECT_EQ(std::count(names.begin(), names.end(), "sdpa_softmax"), s != 1); + EXPECT_EQ( + std::count(names.begin(), names.end(), "sdpa_compute_out"), s != 1); + } +} + +TEST(DynamicShape, StaticRouteProfiles) { + const auto* context = get_default_webgpu_context(); + if (std::getenv("WEBGPU_TIMESTAMP_QUERY") == nullptr || context == nullptr || + !context->timestamp_supported) { + GTEST_SKIP() << "timestamp queries unavailable"; + } + + Module linear_m1(g_dir + "/static_linear_m1.pte"); + ASSERT_EQ(linear_m1.load_forward(), Error::Ok); + run_linear(linear_m1, 1, "static_linear_m1", kLinN); + auto names = current_profile_names(); + ASSERT_EQ(names.size(), 1); + EXPECT_EQ( + std::count(names.begin(), names.end(), "linear_q4gsw_coop4_bicol"), 1); + + Module linear_m32(g_dir + "/static_linear_m32.pte"); + ASSERT_EQ(linear_m32.load_forward(), Error::Ok); + run_linear(linear_m32, 32, "static_linear_m32", kLinN); + names = current_profile_names(); + ASSERT_EQ(names.size(), 1); + EXPECT_EQ(std::count(names.begin(), names.end(), "linear_q4gsw_steel"), 1); + + Module linear_shmem(g_dir + "/dyn_linear_shmem.pte"); + ASSERT_EQ(linear_shmem.load_forward(), Error::Ok); + run_linear(linear_shmem, 128, "dyn_linear_shmem", kLinNShmem, kLinAltK); + names = current_profile_names(); + ASSERT_EQ(names.size(), 1); + EXPECT_EQ(std::count(names.begin(), names.end(), "linear_q4gsw_shmem"), 1); + + Module linear_tiled(g_dir + "/dyn_linear_tiled.pte"); + ASSERT_EQ(linear_tiled.load_forward(), Error::Ok); + run_linear(linear_tiled, 128, "dyn_linear_tiled", kLinN, kLinAltK); + names = current_profile_names(); + ASSERT_EQ(names.size(), 1); + EXPECT_EQ(std::count(names.begin(), names.end(), "linear_q4gsw_tiled"), 1); + + Module sdpa_s1(g_dir + "/static_sdpa_s1.pte"); + ASSERT_EQ(sdpa_s1.load_forward(), Error::Ok); + run_sdpa_case(sdpa_s1, 1, "static_sdpa_s1", kSdHq, kSdHkv, kSdD, kSdCmax); + names = current_profile_names(); + ASSERT_EQ(names.size(), 4); + EXPECT_EQ(std::count(names.begin(), names.end(), "fd_split"), 1); + EXPECT_EQ(std::count(names.begin(), names.end(), "fd_reduce"), 1); + + Module sdpa_s16(g_dir + "/static_sdpa_s16.pte"); + ASSERT_EQ(sdpa_s16.load_forward(), Error::Ok); + run_sdpa_case(sdpa_s16, 16, "static_sdpa_s16", kSdHq, kSdHkv, kSdD, kSdCmax); + names = current_profile_names(); + ASSERT_EQ(names.size(), 5); + EXPECT_EQ( + std::count(names.begin(), names.end(), "sdpa_compute_attn_weights"), 1); + EXPECT_EQ(std::count(names.begin(), names.end(), "sdpa_softmax"), 1); + EXPECT_EQ(std::count(names.begin(), names.end(), "sdpa_compute_out"), 1); +} +#endif + +// J: dynamic SDPA reuses one graph across prefill and FlashDecoding shapes. +TEST(DynamicShape, Sdpa) { for (int s : {64, 16, 1}) { check_sdpa(s); } } +TEST(DynamicShape, SdpaReusedGraph) { + Module m(g_dir + "/sdpa_dyn.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "load sdpa_dyn.pte"; + for (int s : {64, 1, 16, 1, 64}) { + run_sdpa(m, s); + } +} + +TEST(DynamicShape, CombinedLiveRoutes) { + Module m(g_dir + "/combined_routes.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "load combined_routes.pte"; + for (int s : {64, 1, 16, 1, 64}) { + run_combined_routes(m, s); + } +} + +TEST(DynamicShape, SdpaWideMaterializedOnly) { + Module m(g_dir + "/sdpa_wide.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "load sdpa_wide.pte"; + for (int s : {16, 1, 16}) { + run_sdpa_case(m, s, "sdpa_wide", 8, 2, 132, 16); + } +} + +TEST(DynamicShape, K16CausalNumericsReusedGraph) { + if (!k16_device_supported()) { + GTEST_SKIP() << "K16 device limits unavailable"; + } + Module module(g_dir + "/sdpa_k16_llama.pte"); + load_sdpa_module(module, true); + prime_k16_sdpa(module, "sdpa_k16_llama"); + for (int s : {512, 1, 508, 128, 127, 16, 1, 512}) { + run_k16_sdpa(module, s, "sdpa_k16_llama"); + } +} + +TEST(DynamicShape, Qwen3K16CausalNumericsReusedGraph) { + if (!qwen3_q16_device_supported()) { + GTEST_SKIP() << "Qwen3 Q16 K16 device limits unavailable"; + } + constexpr float kQwen3MaxError = 1e-2f; + Module module(g_dir + "/sdpa_k16_qwen3.pte"); + load_sdpa_module(module, true); + prime_k16_sdpa( + module, "sdpa_k16_qwen3", kQwen3Hq, kQwen3Hkv, kQwen3D, kQwen3MaxError); + for (int s : {128, 1, 17, 1, 128}) { + run_k16_sdpa( + module, + s, + "sdpa_k16_qwen3", + kQwen3Hq, + kQwen3Hkv, + kQwen3D, + false, + kQwen3MaxError); + } +} + +TEST(DynamicShape, Qwen3InitializedConstantCachePreserved) { + if (!qwen3_q16_device_supported()) { + GTEST_SKIP() << "Qwen3 Q16 K16 device limits unavailable"; + } + Module module(g_dir + "/sdpa_k16_qwen3.pte"); + load_sdpa_module(module, true); + run_k16_sdpa( + module, + 17, + "sdpa_k16_qwen3", + kQwen3Hq, + kQwen3Hkv, + kQwen3D, + false, + 1e-2f, + true); +} + +TEST(DynamicShape, K16CacheHeadMismatchRejectedAtLoad) { + executorch::runtime::BackendOptions<1> options; + ASSERT_EQ(options.set_option("enable_f16_kv_cache", true), Error::Ok); + executorch::runtime::LoadBackendOptionsMap option_map; + ASSERT_EQ(option_map.set_options("VulkanBackend", options.view()), Error::Ok); + Module module(g_dir + "/sdpa_k16_bad_cache_heads.pte"); + EXPECT_NE(module.load_forward(nullptr, nullptr, &option_map), Error::Ok); +} + +#ifdef WGPU_BACKEND_ENABLE_PROFILING +TEST(DynamicShape, SdpaLiveRoutesProfile) { + const auto* context = get_default_webgpu_context(); + if (std::getenv("WEBGPU_TIMESTAMP_QUERY") == nullptr || context == nullptr || + !context->timestamp_supported) { + GTEST_SKIP() << "timestamp queries unavailable"; + } + Module m(g_dir + "/sdpa_dyn.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "load sdpa_dyn.pte"; + for (int s : {64, 1, 16, 1, 64}) { + run_sdpa(m, s); + const auto names = current_profile_names(); + ASSERT_EQ(names.size(), s == 1 ? 4 : 5); + EXPECT_EQ(std::count(names.begin(), names.end(), "update_cache"), 2); + EXPECT_EQ(std::count(names.begin(), names.end(), "fd_split"), s == 1); + EXPECT_EQ(std::count(names.begin(), names.end(), "fd_reduce"), s == 1); + EXPECT_EQ( + std::count(names.begin(), names.end(), "sdpa_compute_attn_weights"), + s != 1); + EXPECT_EQ(std::count(names.begin(), names.end(), "sdpa_softmax"), s != 1); + EXPECT_EQ( + std::count(names.begin(), names.end(), "sdpa_compute_out"), s != 1); + } +} + +TEST(DynamicShape, K16CausalLiveRoutesProfile) { + const auto* context = get_default_webgpu_context(); + if (std::getenv("WEBGPU_TIMESTAMP_QUERY") == nullptr || context == nullptr || + !context->timestamp_supported || !k16_device_supported()) { + GTEST_SKIP() << "timestamp queries or K16 device limits unavailable"; + } + Module module(g_dir + "/sdpa_k16_llama.pte"); + load_sdpa_module(module, true); + prime_k16_sdpa(module, "sdpa_k16_llama"); + expect_sdpa_route(current_profile_names(), 12, true); + for (int s : {512, 128, 1, 508, 127, 1, 512}) { + run_k16_sdpa(module, s, "sdpa_k16_llama"); + expect_sdpa_route(current_profile_names(), s, s > 1); + } +} + +TEST(DynamicShape, Qwen3K16CausalLiveRoutesProfile) { + const auto* context = get_default_webgpu_context(); + if (std::getenv("WEBGPU_TIMESTAMP_QUERY") == nullptr || context == nullptr || + !context->timestamp_supported || !qwen3_q16_device_supported()) { + GTEST_SKIP() << "timestamp queries or Qwen3 Q16 K16 limits unavailable"; + } + constexpr float kQwen3MaxError = 1e-2f; + constexpr const char* kQwen3Kernel = + "sdpa_streaming_attention_qwen3_k16_causal_bound"; + Module module(g_dir + "/sdpa_k16_qwen3.pte"); + load_sdpa_module(module, true); + prime_k16_sdpa( + module, "sdpa_k16_qwen3", kQwen3Hq, kQwen3Hkv, kQwen3D, kQwen3MaxError); + expect_sdpa_route(current_profile_names(), 12, true, kQwen3Kernel); + for (int s : {128, 1, 17, 1, 128}) { + run_k16_sdpa( + module, + s, + "sdpa_k16_qwen3", + kQwen3Hq, + kQwen3Hkv, + kQwen3D, + false, + kQwen3MaxError); + expect_sdpa_route(current_profile_names(), s, s > 1, kQwen3Kernel); + } +} + +TEST(DynamicShape, Qwen3NearScaleFallsBackToExistingRoutes) { + const auto* context = get_default_webgpu_context(); + if (std::getenv("WEBGPU_TIMESTAMP_QUERY") == nullptr || context == nullptr || + !context->timestamp_supported || !qwen3_q16_device_supported()) { + GTEST_SKIP() << "timestamp queries or Qwen3 Q16 K16 limits unavailable"; + } + constexpr float kQwen3MaxError = 1e-2f; + Module module(g_dir + "/sdpa_k16_qwen3_near_scale.pte"); + load_sdpa_module(module, true); + prime_k16_sdpa( + module, + "sdpa_k16_qwen3_near_scale", + kQwen3Hq, + kQwen3Hkv, + kQwen3D, + kQwen3MaxError); + expect_sdpa_route(current_profile_names(), 12, false); + for (int s : {128, 1, 128}) { + run_k16_sdpa( + module, + s, + "sdpa_k16_qwen3_near_scale", + kQwen3Hq, + kQwen3Hkv, + kQwen3D, + false, + kQwen3MaxError); + expect_sdpa_route(current_profile_names(), s, false); + } +} + +TEST(DynamicShape, K16F32KvFallsBackToExistingRoutes) { + const auto* context = get_default_webgpu_context(); + if (std::getenv("WEBGPU_TIMESTAMP_QUERY") == nullptr || context == nullptr || + !context->timestamp_supported || !k16_device_supported()) { + GTEST_SKIP() << "timestamp queries or K16 device limits unavailable"; + } + Module module(g_dir + "/sdpa_k16_llama.pte"); + load_sdpa_module(module, false); + prime_k16_sdpa(module, "sdpa_k16_llama"); + expect_sdpa_route(current_profile_names(), 12, false); + for (int s : {128, 1, 128}) { + run_k16_sdpa(module, s, "sdpa_k16_llama"); + expect_sdpa_route(current_profile_names(), s, false); + } +} + +TEST(DynamicShape, K16MetadataFallsBackToExistingRoutes) { + const auto* context = get_default_webgpu_context(); + if (std::getenv("WEBGPU_TIMESTAMP_QUERY") == nullptr || context == nullptr || + !context->timestamp_supported || !k16_device_supported()) { + GTEST_SKIP() << "timestamp queries or K16 device limits unavailable"; + } + struct NegativeCase { + const char* prefix; + int hq; + int hkv; + int d; + }; + for (const auto& negative : std::vector{ + {"sdpa_k16_wrong_geometry", 14, 2, 64}, + {"sdpa_k16_wrong_d", 32, 8, 128}, + {"sdpa_k16_wrong_scale", 32, 8, 64}}) { + Module module(g_dir + "/" + negative.prefix + ".pte"); + load_sdpa_module(module, true); + prime_k16_sdpa( + module, negative.prefix, negative.hq, negative.hkv, negative.d); + expect_sdpa_route(current_profile_names(), 12, false); + for (int s : {128, 1, 128}) { + run_k16_sdpa( + module, s, negative.prefix, negative.hq, negative.hkv, negative.d); + expect_sdpa_route(current_profile_names(), s, false); + } + } +} +#endif + // K: dynamic embedding (int64 token ids) at several token counts. TEST(DynamicShape, Embedding) { for (int n : {16, 8, 1}) { @@ -459,12 +1696,14 @@ TEST(DynamicShape, EmbeddingReusedGraph) { } } -// K3: linear-packed reuse must preserve nibble order across resizes. -TEST(DynamicShape, LinearPackedEmbeddingReusedGraph) { - Module m(g_dir + "/emb_dyn_linear.pte"); - ASSERT_EQ(m.load_forward(), Error::Ok) << "load emb_dyn_linear.pte"; - for (int n : {16, 8, 1, 16}) { - run_embedding(m, n, "emb_dyn_linear"); +// K3: linear/nonlinear-packed reuse must preserve nibble order across resizes. +TEST(DynamicShape, EmbeddingLayoutsReusedGraph) { + for (const char* prefix : {"emb_dyn_linear", "emb_dyn_nonlinear"}) { + Module m(g_dir + "/" + prefix + ".pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "load " << prefix << ".pte"; + for (int n : {16, 8, 1, 16}) { + run_embedding(m, n, prefix); + } } } @@ -503,6 +1742,224 @@ TEST(DynamicShape, SigmoidReusedGraph) { } } +TEST(DynamicShape, SwiGluReusedGraph) { + Module module(g_dir + "/dyn_swiglu.pte"); + ASSERT_EQ(module.load_forward(), Error::Ok) << "load dyn_swiglu.pte"; + for (int m_rows : {512, 128, 1, 512}) { + run_swiglu(module, m_rows, "dyn_swiglu", kSwiGluWidth); + } +} + +TEST(DynamicShape, SwiGluCommutativeAndOwnership) { + for (const char* prefix : + {"dyn_swiglu_inner_reversed", "dyn_swiglu_outer_reversed"}) { + Module module(g_dir + "/" + prefix + ".pte"); + ASSERT_EQ(module.load_forward(), Error::Ok) << "load " << prefix; + run_swiglu(module, 128, prefix, kSwiGluSmallWidth); + } + + Module negative(g_dir + "/dyn_swiglu_extra_gate_consumer.pte"); + ASSERT_EQ(negative.load_forward(), Error::Ok) + << "load dyn_swiglu_extra_gate_consumer"; + run_swiglu( + negative, 128, "dyn_swiglu_extra_gate_consumer", kSwiGluSmallWidth); + + Module graph_outputs(g_dir + "/dyn_swiglu_graph_outputs.pte"); + ASSERT_EQ(graph_outputs.load_forward(), Error::Ok) + << "load dyn_swiglu_graph_outputs"; + run_swiglu_graph_outputs(graph_outputs, 128); + + for (const char* prefix : + {"dyn_swiglu_extra_sigmoid_consumer", + "dyn_swiglu_extra_silu_consumer"}) { + Module module(g_dir + "/" + prefix + ".pte"); + ASSERT_EQ(module.load_forward(), Error::Ok) << "load " << prefix; + run_swiglu(module, 128, prefix, kSwiGluSmallWidth); + } + + for (const char* prefix : + {"dyn_swiglu_gate_graph_output", + "dyn_swiglu_sigmoid_graph_output", + "dyn_swiglu_silu_graph_output"}) { + Module module(g_dir + "/" + prefix + ".pte"); + ASSERT_EQ(module.load_forward(), Error::Ok) << "load " << prefix; + run_swiglu_outputs(module, 128, prefix, 2); + } + + Module different_inputs(g_dir + "/dyn_swiglu_different_inputs.pte"); + ASSERT_EQ(different_inputs.load_forward(), Error::Ok); + run_swiglu( + different_inputs, + 128, + "dyn_swiglu_different_inputs", + kSwiGluSmallWidth, + true); + + Module interleaved(g_dir + "/dyn_swiglu_interleaved_q4.pte"); + ASSERT_EQ(interleaved.load_forward(), Error::Ok); + run_swiglu(interleaved, 128, "dyn_swiglu_interleaved_q4", kSwiGluSmallWidth); + + Module overlap(g_dir + "/dyn_swiglu_qkv_overlap.pte"); + ASSERT_EQ(overlap.load_forward(), Error::Ok); + run_swiglu_qkv_overlap(overlap, 128); +} + +#ifdef WGPU_BACKEND_ENABLE_PROFILING +void expect_swiglu_profile( + Module& module, + int m_rows, + const char* prefix, + int width, + bool expect_2d) { + run_swiglu(module, m_rows, prefix, width); + const auto* context = get_default_webgpu_context(); + ASSERT_NE(context, nullptr); + ASSERT_NE(context->querypool, nullptr); + const auto& profile = context->querypool->results(); + ASSERT_EQ(profile.size(), 3) + << "two q4 projections plus one fused SwiGLU dispatch expected"; + EXPECT_EQ( + std::count_if( + profile.begin(), + profile.end(), + [](const auto& duration) { + return duration.kernel_name == "silu_mul_fused"; + }), + 1); + EXPECT_EQ( + std::count_if( + profile.begin(), + profile.end(), + [](const auto& duration) { + return duration.kernel_name == "mul" || + duration.kernel_name == "sigmoid"; + }), + 0); + const auto fused = + std::find_if(profile.begin(), profile.end(), [](const auto& duration) { + return duration.kernel_name == "silu_mul_fused"; + }); + ASSERT_NE(fused, profile.end()); + EXPECT_EQ(fused->global_wg[1] > 1, expect_2d); +} + +TEST(DynamicShape, SwiGluFusionProfile) { + const auto* context = get_default_webgpu_context(); + if (std::getenv("WEBGPU_TIMESTAMP_QUERY") == nullptr || context == nullptr || + !context->timestamp_supported) { + GTEST_SKIP() << "timestamp queries unavailable"; + } + + for (const char* prefix : + {"dyn_swiglu_inner_reversed", "dyn_swiglu_outer_reversed"}) { + Module module(g_dir + "/" + prefix + ".pte"); + ASSERT_EQ(module.load_forward(), Error::Ok) << prefix; + expect_swiglu_profile(module, 128, prefix, kSwiGluSmallWidth, false); + } + + Module canonical(g_dir + "/dyn_swiglu.pte"); + ASSERT_EQ(canonical.load_forward(), Error::Ok); + for (int m_rows : {1, 128, 512}) { + expect_swiglu_profile( + canonical, m_rows, "dyn_swiglu", kSwiGluWidth, m_rows == 512); + } + + Module negative(g_dir + "/dyn_swiglu_extra_gate_consumer.pte"); + ASSERT_EQ(negative.load_forward(), Error::Ok); + run_swiglu( + negative, 128, "dyn_swiglu_extra_gate_consumer", kSwiGluSmallWidth); + const auto names = current_profile_names(); + EXPECT_FALSE(contains_name(names, "silu_mul_fused")); + EXPECT_EQ(std::count(names.begin(), names.end(), "mul"), 2); + + Module graph_outputs(g_dir + "/dyn_swiglu_graph_outputs.pte"); + ASSERT_EQ(graph_outputs.load_forward(), Error::Ok); + run_swiglu_graph_outputs(graph_outputs, 128); + const auto graph_output_names = current_profile_names(); + EXPECT_FALSE(contains_name(graph_output_names, "silu_mul_fused")); + EXPECT_EQ( + std::count(graph_output_names.begin(), graph_output_names.end(), "mul"), + 2); + + for (const char* prefix : + {"dyn_swiglu_extra_sigmoid_consumer", + "dyn_swiglu_extra_silu_consumer"}) { + Module module(g_dir + "/" + prefix + ".pte"); + ASSERT_EQ(module.load_forward(), Error::Ok) << prefix; + run_swiglu(module, 128, prefix, kSwiGluSmallWidth); + const auto consumer_names = current_profile_names(); + EXPECT_FALSE(contains_name(consumer_names, "silu_mul_fused")) << prefix; + EXPECT_EQ( + std::count(consumer_names.begin(), consumer_names.end(), "mul"), 2) + << prefix; + } + + for (const char* prefix : + {"dyn_swiglu_gate_graph_output", + "dyn_swiglu_sigmoid_graph_output", + "dyn_swiglu_silu_graph_output"}) { + Module module(g_dir + "/" + prefix + ".pte"); + ASSERT_EQ(module.load_forward(), Error::Ok) << prefix; + run_swiglu_outputs(module, 128, prefix, 2); + const auto output_names = current_profile_names(); + EXPECT_FALSE(contains_name(output_names, "silu_mul_fused")) << prefix; + EXPECT_EQ(std::count(output_names.begin(), output_names.end(), "mul"), 2) + << prefix; + } + + Module different_inputs(g_dir + "/dyn_swiglu_different_inputs.pte"); + ASSERT_EQ(different_inputs.load_forward(), Error::Ok); + run_swiglu( + different_inputs, + 128, + "dyn_swiglu_different_inputs", + kSwiGluSmallWidth, + true); + const auto different_input_names = current_profile_names(); + EXPECT_FALSE(contains_name(different_input_names, "silu_mul_fused")); + EXPECT_EQ( + std::count( + different_input_names.begin(), different_input_names.end(), "mul"), + 2); + + Module interleaved(g_dir + "/dyn_swiglu_interleaved_q4.pte"); + ASSERT_EQ(interleaved.load_forward(), Error::Ok); + run_swiglu(interleaved, 128, "dyn_swiglu_interleaved_q4", kSwiGluSmallWidth); + const auto interleaved_names = current_profile_names(); + EXPECT_EQ(interleaved_names.size(), 5); + EXPECT_EQ( + std::count( + interleaved_names.begin(), interleaved_names.end(), "silu_mul_fused"), + 1); + EXPECT_EQ( + std::count(interleaved_names.begin(), interleaved_names.end(), "mul"), 0); +} + +TEST(DynamicShape, SwiGluQkvOverlapProfile) { + const auto* context = get_default_webgpu_context(); + if (std::getenv("WEBGPU_TIMESTAMP_QUERY") == nullptr || context == nullptr || + !context->timestamp_supported || !context->shader_f16_supported) { + GTEST_SKIP() << "timestamp queries or shader-f16 unavailable"; + } + WGPULimits limits = {}; + if (wgpuDeviceGetLimits(context->device, &limits) != WGPUStatus_Success || + limits.maxComputeInvocationsPerWorkgroup < 256u || + limits.maxComputeWorkgroupSizeX < 16u || + limits.maxComputeWorkgroupSizeY < 16u || + limits.maxComputeWorkgroupStorageSize < 16384u || + limits.maxComputeWorkgroupsPerDimension < 384u) { + GTEST_SKIP() << "QKV workgroup limits unavailable"; + } + Module overlap(g_dir + "/dyn_swiglu_qkv_overlap.pte"); + ASSERT_EQ(overlap.load_forward(), Error::Ok); + run_swiglu_qkv_overlap(overlap, 128); + const auto names = current_profile_names(); + EXPECT_EQ(std::count(names.begin(), names.end(), "linear_q4gsw_bk64_qkv"), 0); + EXPECT_EQ(std::count(names.begin(), names.end(), "silu_mul_fused"), 1); + EXPECT_EQ(std::count(names.begin(), names.end(), "mul"), 0); +} +#endif + // N: dynamic select_copy(0,-1) at several S. TEST(DynamicShape, Select) { for (int s : {128, 32, 1}) { diff --git a/backends/webgpu/test/native/test_execution_options.cpp b/backends/webgpu/test/native/test_execution_options.cpp new file mode 100644 index 00000000000..382bfbf7efa --- /dev/null +++ b/backends/webgpu/test/native/test_execution_options.cpp @@ -0,0 +1,167 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include + +#include + +#include +#include + +namespace executorch::backends::webgpu { +namespace { + +TEST(WebGPUExecutionOptionsTest, DefaultsToPreservingOutputs) { + EXPECT_EQ( + current_webgpu_execution_options().discardable_output_data, nullptr); +} + +TEST(WebGPUExecutionOptionsTest, NestedScopesRestorePriorValue) { + int outer_output = 0; + int inner_output = 0; + { + ScopedWebGPUExecutionOptions outer({&outer_output, true}); + EXPECT_EQ( + current_webgpu_execution_options().discardable_output_data, + &outer_output); + { + ScopedWebGPUExecutionOptions inner({&inner_output, true}); + EXPECT_EQ( + current_webgpu_execution_options().discardable_output_data, + &inner_output); + } + EXPECT_EQ( + current_webgpu_execution_options().discardable_output_data, + &outer_output); + } + EXPECT_EQ( + current_webgpu_execution_options().discardable_output_data, nullptr); +} + +TEST(WebGPUExecutionOptionsTest, ExceptionRestoresPriorValue) { + int output = 0; + EXPECT_THROW( + with_webgpu_execution_options( + {&output, true}, + []() -> void { throw std::runtime_error("expected"); }), + std::runtime_error); + EXPECT_EQ( + current_webgpu_execution_options().discardable_output_data, nullptr); +} + +TEST(WebGPUExecutionOptionsTest, BooleanReturnRestoresPriorValue) { + int output = 0; + const bool result = + with_webgpu_execution_options({&output, true}, []() { return false; }); + EXPECT_FALSE(result); + EXPECT_EQ( + current_webgpu_execution_options().discardable_output_data, nullptr); +} + +TEST(WebGPUExecutionOptionsTest, ResolvesOnlyOneExactDelegateOutput) { + int method_output = 0; + int delegate_intermediate = 0; + const std::vector delegate_outputs = {&delegate_intermediate}; + + EXPECT_EQ( + resolve_webgpu_graph_execution_options( + delegate_outputs, WebGPUExecutionOptions{&method_output, true}) + .suppress_output_ordinal, + kNoOutputOrdinal); + EXPECT_EQ( + resolve_webgpu_graph_execution_options( + {&delegate_intermediate, &method_output}, + WebGPUExecutionOptions{&method_output, true}) + .suppress_output_ordinal, + 1); + EXPECT_EQ( + resolve_webgpu_graph_execution_options( + {&method_output, &method_output}, + WebGPUExecutionOptions{&method_output, true}) + .suppress_output_ordinal, + kNoOutputOrdinal); + EXPECT_EQ( + resolve_webgpu_graph_execution_options( + {&method_output}, WebGPUExecutionOptions{&method_output, false}) + .suppress_output_ordinal, + kNoOutputOrdinal); +} + +TEST(WebGPUExecutionPlanTest, DefaultPlanPreservesDispatchesAndOutputs) { + const std::vector suppressible = {{9, 1, 4, 6}}; + const WebGPUExecutionPlan plan = plan_webgpu_execution( + 6, 2, ExecuteConfig{}, suppressible, WebGPUGraphExecutionOptions{}); + + EXPECT_EQ( + plan.dispatch_chunks, + (std::vector>{{0, 1, 2, 3, 4, 5}})); + EXPECT_EQ(plan.copy_outputs, (std::vector{true, true})); +} + +TEST(WebGPUExecutionPlanTest, SuppressionIsPerOutputAndSupportsChunking) { + const std::vector suppressible = {{9, 1, 4, 6}}; + const ExecuteConfig config = {2, 1}; + const WebGPUExecutionPlan plan = plan_webgpu_execution( + 6, 2, config, suppressible, WebGPUGraphExecutionOptions{1}); + + EXPECT_EQ( + plan.dispatch_chunks, + (std::vector>{{0}, {1, 2}, {3}})); + EXPECT_EQ(plan.copy_outputs, (std::vector{true, false})); +} + +TEST(WebGPUExecutionPlanTest, RejectsInvalidSuppressibleRange) { + const std::vector suppressible = {{9, 0, 3, 7}}; + EXPECT_THROW( + plan_webgpu_execution( + 6, 1, ExecuteConfig{}, suppressible, WebGPUGraphExecutionOptions{0}), + std::runtime_error); +} + +TEST(WebGPUExecutionPlanTest, AllSuppressedHasNoSyntheticDispatchChunk) { + const std::vector suppressible = {{9, 0, 0, 2}}; + const WebGPUExecutionPlan plan = plan_webgpu_execution( + 2, 1, ExecuteConfig{}, suppressible, WebGPUGraphExecutionOptions{0}); + + EXPECT_TRUE(plan.dispatch_chunks.empty()); + EXPECT_EQ(plan.copy_outputs, (std::vector{false})); +} + +TEST(WebGPUExecutionPlanTest, CopyOnlyPlanRetainsOneSubmissionChunk) { + const WebGPUExecutionPlan plan = plan_webgpu_execution( + 0, 1, ExecuteConfig{}, {}, WebGPUGraphExecutionOptions{}); + + EXPECT_EQ(plan.dispatch_chunks, (std::vector>{{}})); + EXPECT_EQ(plan.copy_outputs, (std::vector{true})); +} + +TEST(WebGPUExecutionPlanTest, FiltersDisabledDispatchesAcrossChunks) { + const std::vector enabled = {true, false, true, false, true, true}; + const WebGPUExecutionPlan plan = plan_webgpu_execution( + 6, 1, ExecuteConfig{2, 1}, {}, WebGPUGraphExecutionOptions{}, enabled); + + EXPECT_EQ( + plan.dispatch_chunks, + (std::vector>{{0}, {2}, {4}, {5}})); + EXPECT_EQ(plan.copy_outputs, (std::vector{true})); +} + +TEST(WebGPUExecutionPlanTest, RejectsMismatchedEnabledDispatches) { + EXPECT_THROW( + plan_webgpu_execution( + 3, + 1, + ExecuteConfig{}, + {}, + WebGPUGraphExecutionOptions{}, + {true, false}), + std::runtime_error); +} + +} // namespace +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/test/native/test_output_suppression.cpp b/backends/webgpu/test/native/test_output_suppression.cpp new file mode 100644 index 00000000000..e5012047caf --- /dev/null +++ b/backends/webgpu/test/native/test_output_suppression.cpp @@ -0,0 +1,180 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace executorch::backends::webgpu; +using namespace executorch::extension; +using namespace executorch::runtime; + +namespace { + +constexpr int kWidth = 64; +constexpr float kPoison = -12345.0f; + +std::string g_dir; // NOLINT(cppcoreguidelines-avoid-non-const-global-variables) + +std::vector read_bin(const std::string& path) { + std::ifstream file(path, std::ios::binary | std::ios::ate); + if (!file) { + fprintf(stderr, "read_bin: cannot open %s\n", path.c_str()); + return {}; + } + const std::streamsize bytes = file.tellg(); + if (bytes < 0 || bytes % sizeof(float) != 0) { + fprintf( + stderr, + "read_bin: %s has non-float size %lld\n", + path.c_str(), + static_cast(bytes)); + return {}; + } + file.seekg(0); + std::vector values(static_cast(bytes) / sizeof(float)); + file.read(reinterpret_cast(values.data()), bytes); + return values; +} + +void expect_matches( + const std::vector& actual, + const std::vector& expected) { + ASSERT_EQ(actual.size(), expected.size()); + for (size_t i = 0; i < actual.size(); i++) { + const float abs_err = std::fabs(actual[i] - expected[i]); + const float denom = std::max(std::fabs(expected[i]), 1e-12f); + const float rel_err = abs_err / denom; + EXPECT_TRUE(abs_err <= 1e-3f || rel_err <= 1e-3f) + << "index=" << i << " actual=" << actual[i] + << " expected=" << expected[i] << " abs=" << abs_err + << " rel=" << rel_err; + } +} + +std::vector input_values() { + auto input = read_bin(g_dir + "/input.bin"); + EXPECT_EQ(input.size(), static_cast(kWidth)); + return input; +} + +Result> forward_with_suppression( + Module& module, + const TensorPtr& input, + const void* discardable_output_data) { + return with_webgpu_execution_options( + WebGPUExecutionOptions{discardable_output_data, true}, + [&]() { return module.forward({EValue(input)}); }); +} + +TEST(OutputSuppression, DirectFinalQ4SkipsOnlyDiscardedInvocation) { + Module module(g_dir + "/direct_final_q4.pte"); + ASSERT_EQ(module.load_forward(), Error::Ok); + + auto input = make_tensor_ptr({1, kWidth}, input_values()); + std::vector suppressed(kWidth, kPoison); + auto suppressed_tensor = + make_tensor_ptr({1, kWidth}, static_cast(suppressed.data())); + ASSERT_EQ(module.set_output(EValue(suppressed_tensor)), Error::Ok); + const auto result = + forward_with_suppression(module, input, suppressed.data()); + ASSERT_TRUE(result.ok()); + EXPECT_TRUE(std::all_of(suppressed.begin(), suppressed.end(), [](float v) { + return v == kPoison; + })); + + std::vector terminal(kWidth, kPoison); + auto terminal_tensor = + make_tensor_ptr({1, kWidth}, static_cast(terminal.data())); + ASSERT_EQ(module.set_output(EValue(terminal_tensor)), Error::Ok); + const auto terminal_result = module.forward({EValue(input)}); + ASSERT_TRUE(terminal_result.ok()); + expect_matches( + terminal, read_bin(g_dir + "/direct_final_q4.output0.golden.bin")); +} + +TEST(OutputSuppression, Q4FeedingLaterAddIsNotSuppressible) { + Module module(g_dir + "/q4_then_add.pte"); + ASSERT_EQ(module.load_forward(), Error::Ok); + + auto input = make_tensor_ptr({1, kWidth}, input_values()); + std::vector output0(kWidth, kPoison); + std::vector output1(kWidth, kPoison); + auto tensor0 = + make_tensor_ptr({1, kWidth}, static_cast(output0.data())); + auto tensor1 = + make_tensor_ptr({1, kWidth}, static_cast(output1.data())); + ASSERT_EQ(module.set_outputs({EValue(tensor0), EValue(tensor1)}), Error::Ok); + const auto result = forward_with_suppression(module, input, output0.data()); + ASSERT_TRUE(result.ok()); + expect_matches(output0, read_bin(g_dir + "/q4_then_add.output0.golden.bin")); + expect_matches(output1, read_bin(g_dir + "/q4_then_add.output1.golden.bin")); +} + +TEST(OutputSuppression, UnrelatedOutputStillCopies) { + Module module(g_dir + "/unrelated_then_final_q4.pte"); + ASSERT_EQ(module.load_forward(), Error::Ok); + + auto input = make_tensor_ptr({1, kWidth}, input_values()); + std::vector unrelated(kWidth, kPoison); + std::vector suppressed(kWidth, kPoison); + auto unrelated_tensor = + make_tensor_ptr({1, kWidth}, static_cast(unrelated.data())); + auto suppressed_tensor = + make_tensor_ptr({1, kWidth}, static_cast(suppressed.data())); + ASSERT_EQ( + module.set_outputs({EValue(unrelated_tensor), EValue(suppressed_tensor)}), + Error::Ok); + const auto result = + forward_with_suppression(module, input, suppressed.data()); + ASSERT_TRUE(result.ok()); + expect_matches( + unrelated, + read_bin(g_dir + "/unrelated_then_final_q4.output0.golden.bin")); + EXPECT_TRUE(std::all_of(suppressed.begin(), suppressed.end(), [](float v) { + return v == kPoison; + })); +} + +} // namespace + +int main(int argc, char** argv) { + ::testing::InitGoogleTest(&argc, argv); + g_dir = "/tmp/output_suppression"; + if (argc > 1) { + g_dir = argv[1]; + } + if (const char* env = std::getenv("WEBGPU_OUTPUT_SUPPRESSION_DIR")) { + g_dir = env; + } + + WebGPUContext context; + try { + context = create_webgpu_context(); + } catch (const std::exception& error) { + std::printf("SKIP: no WebGPU device (%s)\n", error.what()); + return 0; + } + set_default_webgpu_context(&context); + const int result = RUN_ALL_TESTS(); + set_default_webgpu_context(nullptr); + destroy_webgpu_context(context); + return result; +} diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index 19438d6acbc..2c78e9192ca 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -38,10 +38,10 @@ ) from executorch.backends.webgpu.test.ops.test_avg_pool2d import AvgPool2dModule from executorch.backends.webgpu.test.ops.test_bitwise import ( + BITWISE_NOT_SHAPES, BitwiseAndModule, BitwiseNotModule, bw_gen_a, - bw_gen_b, ) from executorch.backends.webgpu.test.ops.test_cat import ( CatModule, @@ -71,14 +71,13 @@ make_qcs4w_linear_module, ) from executorch.backends.webgpu.test.ops.test_logical_and import ( - la_gen_a, - la_gen_b, + LOGICAL_BINARY_CASES, + logical_binary_gen_a, + logical_binary_gen_b, LogicalAndModule, ) from executorch.backends.webgpu.test.ops.test_logical_or import ( BitwiseOrModule, - lo_gen_a, - lo_gen_b, LogicalOrModule, ) from executorch.backends.webgpu.test.ops.test_minimum import MinimumModule @@ -112,7 +111,15 @@ DequantizeConstModule, QuantizeModule, ) -from executorch.backends.webgpu.test.ops.test_reduce import AmaxModule, AminModule +from executorch.backends.webgpu.test.ops.test_reduce import ( + amax_sign_trap_input, + amax_tie_input, + AmaxModule, + amin_sign_trap_input, + amin_tie_input, + AminModule, + EXTREMA_CONFIGS, +) from executorch.backends.webgpu.test.ops.test_repeat import RepeatModule from executorch.backends.webgpu.test.ops.test_rms_norm import ( _CASES, @@ -145,6 +152,18 @@ SqueezeModule, ) +from executorch.backends.webgpu.test.ops.test_to_copy import ( + bool_tail_input, + compare_to_copy_input_a, + compare_to_copy_input_b, + CompareToCopyBoolToFloatModule, + to_copy_float_input, + to_copy_int_input, + ToCopyBoolToFloatModule, + ToCopyFloatToIntToFloatModule, + ToCopyIntToFloatModule, +) + from executorch.backends.webgpu.test.ops.test_unary_activations import ( _lin as _unary_lin, CLAMP_CONFIGS, @@ -177,6 +196,39 @@ def _add_factory(variant: str = "regular") -> torch.nn.Module: }[variant]() +@register_op_test("to_copy_bool_to_float") +def _to_copy_bool_to_float_suite() -> WebGPUTestSuite: + return WebGPUTestSuite( + module_factory=CompareToCopyBoolToFloatModule, + cases=[ + Case( + inputs=( + InputSpec((n,), gen=compare_to_copy_input_a), + InputSpec((n,), gen=compare_to_copy_input_b), + ), + name=f"length_{n}", + ) + for n in (1, 4, 5, 67) + ], + golden_dtype="float32", + ) + + +@register_op_test("to_copy_bool_input_to_float") +def _to_copy_bool_input_to_float_suite() -> WebGPUTestSuite: + return WebGPUTestSuite( + module_factory=ToCopyBoolToFloatModule, + cases=[ + Case( + inputs=(InputSpec((n,), gen=bool_tail_input),), + name=f"length_{n}", + ) + for n in (1, 4, 5, 67) + ], + golden_dtype="float32", + ) + + @register_op_test("add") def _add_suite() -> WebGPUTestSuite: # Same-shape numeric coverage only: broadcast adds stay export-smoke in @@ -238,7 +290,7 @@ def _rms_norm_suite() -> WebGPUTestSuite: @register_op_test("mul") def _mul_suite() -> WebGPUTestSuite: - # Full numeric coverage incl. broadcast (binary_mul.wgsl over a TensorMeta UBO); fp64 golden. + # Full binary_op-family numeric coverage, including broadcast. return WebGPUTestSuite( module_factory=lambda: MulModule(), cases=[ @@ -263,21 +315,25 @@ def _fn_config_suite(module_cls, configs) -> WebGPUTestSuite: @register_op_test("minimum") def _minimum_suite() -> WebGPUTestSuite: - # Same-shape numeric coverage (flat binary kernel; broadcast stays smoke). + # Same-shape and mixed-rank broadcast numeric coverage. return WebGPUTestSuite( module_factory=lambda: MinimumModule(), cases=[ Case(name="2d", inputs=((M1, M2), (M1, M2))), Case(name="3d", inputs=((S, S1, S2), (S, S1, S2))), + Case( + name="broadcast_3d_2d", + inputs=( + InputSpec(shape=(2, 3, 8), gen=_unary_lin(-3.0, 3.0)), + InputSpec(shape=(3, 1), gen=_unary_lin(-2.0, 4.0)), + ), + ), ], ) def _compare_suite(op: str) -> WebGPUTestSuite: - # Elementwise fp32 comparison -> bool (byte-exact golden). The two inputs use - # DIFFERENT discrete-range seeds so a!=b (real lt/gt mix) while colliding - # often (eq/le/ge ties); all shapes have numel % 4 == 0 (bool output packs 4 - # bytes/word). Same-shape only (flat kernel; broadcast=smoke). + # Distinct inputs and tail shapes cover byte-exact packed BOOL output. def case(name, shape): return Case( name=name, @@ -289,7 +345,14 @@ def case(name, shape): return WebGPUTestSuite( module_factory=lambda: CompareModule(op), - cases=[case("2d", (4, 8)), case("3d", (2, 3, 8)), case("sq", (16, 16))], + cases=[ + case("tail_1", (1,)), + case("tail_5", (5,)), + case("tail_67", (67,)), + case("2d", (4, 8)), + case("3d", (2, 3, 8)), + case("sq", (16, 16)), + ], golden_dtype="bool", ) @@ -319,49 +382,33 @@ def _ge_suite() -> WebGPUTestSuite: return _compare_suite("ge") -@register_op_test("logical_and") -def _logical_and_suite() -> WebGPUTestSuite: - # out = (a>0) && (b>0): two bool masks derived on-GPU from float inputs via - # gt.Tensor (baked zeros), AND'd -> bool. Distinct a/b seeds so the masks - # differ (AND ~25% True, a real mix an OR mutant fails); all shapes numel % - # 4 == 0 (bool packs 4/word). float32 oracle (byte-exact bool golden). +def _logical_binary_suite(module_factory) -> WebGPUTestSuite: + # Every packed u32 receives all four Boolean pairs in byte order. def case(name, shape): return Case( name=name, construct={"shape": shape}, inputs=( - InputSpec(shape=shape, gen=la_gen_a), - InputSpec(shape=shape, gen=la_gen_b), + InputSpec(shape=shape, gen=logical_binary_gen_a), + InputSpec(shape=shape, gen=logical_binary_gen_b), ), ) return WebGPUTestSuite( - module_factory=lambda shape: LogicalAndModule(shape), - cases=[case("2d", (4, 8)), case("3d", (2, 3, 8)), case("sq", (16, 16))], + module_factory=module_factory, + cases=[case(name, shape) for name, shape in LOGICAL_BINARY_CASES], golden_dtype="float32", ) +@register_op_test("logical_and") +def _logical_and_suite() -> WebGPUTestSuite: + return _logical_binary_suite(lambda shape: LogicalAndModule(shape)) + + @register_op_test("bitwise_and") def _bitwise_and_suite() -> WebGPUTestSuite: - # bool bitwise AND == logical_and for canonical 0/1 (shares the handler). - # Two masks derived on-GPU from float inputs via gt.Tensor (baked zeros), - # distinct a/b seeds (AND ~25% True); all shapes numel % 4 == 0. - def case(name, shape): - return Case( - name=name, - construct={"shape": shape}, - inputs=( - InputSpec(shape=shape, gen=bw_gen_a), - InputSpec(shape=shape, gen=bw_gen_b), - ), - ) - - return WebGPUTestSuite( - module_factory=lambda shape: BitwiseAndModule(shape), - cases=[case("2d", (4, 8)), case("3d", (2, 3, 8)), case("sq", (16, 16))], - golden_dtype="float32", - ) + return _logical_binary_suite(lambda shape: BitwiseAndModule(shape)) @register_op_test("bitwise_not") @@ -377,52 +424,22 @@ def case(name, shape): return WebGPUTestSuite( module_factory=lambda shape: BitwiseNotModule(shape), - cases=[case("2d", (4, 8)), case("3d", (2, 3, 8)), case("sq", (16, 16))], + cases=[ + case(name, shape) + for name, shape in zip(("2d", "3d", "sq"), BITWISE_NOT_SHAPES) + ], golden_dtype="float32", ) @register_op_test("logical_or") def _logical_or_suite() -> WebGPUTestSuite: - # out = (a>0) || (b>0): two bool masks derived on-GPU from float inputs via - # gt.Tensor (baked zeros), OR'd -> bool. Distinct a/b seeds (~50% each, - # independent -> OR ~75% True, a real mix an AND mutant fails); all shapes - # numel % 4 == 0. float32 oracle (byte-exact bool golden). - def case(name, shape): - return Case( - name=name, - construct={"shape": shape}, - inputs=( - InputSpec(shape=shape, gen=lo_gen_a), - InputSpec(shape=shape, gen=lo_gen_b), - ), - ) - - return WebGPUTestSuite( - module_factory=lambda shape: LogicalOrModule(shape), - cases=[case("2d", (4, 8)), case("3d", (2, 3, 8)), case("sq", (16, 16))], - golden_dtype="float32", - ) + return _logical_binary_suite(lambda shape: LogicalOrModule(shape)) @register_op_test("bitwise_or") def _bitwise_or_suite() -> WebGPUTestSuite: - # bool bitwise OR == logical_or for canonical 0/1 (shares the handler). - def case(name, shape): - return Case( - name=name, - construct={"shape": shape}, - inputs=( - InputSpec(shape=shape, gen=lo_gen_a), - InputSpec(shape=shape, gen=lo_gen_b), - ), - ) - - return WebGPUTestSuite( - module_factory=lambda shape: BitwiseOrModule(shape), - cases=[case("2d", (4, 8)), case("3d", (2, 3, 8)), case("sq", (16, 16))], - golden_dtype="float32", - ) + return _logical_binary_suite(lambda shape: BitwiseOrModule(shape)) @register_op_test("pow") @@ -445,6 +462,13 @@ def _pow_suite() -> WebGPUTestSuite: InputSpec(shape=(S, S1, S2), gen=_unary_lin(-2.0, 3.0)), ), ), + Case( + name="broadcast_3d_2d", + inputs=( + InputSpec(shape=(2, 3, 8), gen=_unary_lin(0.1, 3.0)), + InputSpec(shape=(3, 1), gen=_unary_lin(-2.0, 3.0)), + ), + ), ], ) @@ -478,39 +502,54 @@ def _floor_divide_suite() -> WebGPUTestSuite: ), golden_fn=_floor_div_golden, ), + Case( + name="broadcast_3d_2d", + inputs=( + InputSpec(shape=(2, 3, 8), gen=_unary_lin(-8.0, 8.0)), + InputSpec(shape=(3, 1), gen=_unary_lin(0.5, 4.0)), + ), + golden_fn=_floor_div_golden, + ), ], ) -def _reduce_suite(module_cls) -> WebGPUTestSuite: - # Last-dim reduction; both keepdim variants over a 2d and a 3d shape. - return WebGPUTestSuite( - module_factory=lambda keepdim: module_cls(keepdim), - cases=[ - Case(name="keepdim_2d", construct={"keepdim": True}, inputs=((M1, M2),)), - Case(name="nodim_2d", construct={"keepdim": False}, inputs=((M1, M2),)), - Case( - name="keepdim_3d", - construct={"keepdim": True}, - inputs=((S, S1, S2),), - ), +def _reduce_suite(module_cls, op: str) -> WebGPUTestSuite: + generators = { + ("amax", "sign_trap"): amax_sign_trap_input, + ("amax", "tie"): amax_tie_input, + ("amin", "sign_trap"): amin_sign_trap_input, + ("amin", "tie"): amin_tie_input, + } + cases = [] + for name, shape, dim, keepdim, input_class in EXTREMA_CONFIGS: + inputs = (shape,) + kwargs = {} + if input_class != "default": + inputs = (InputSpec(shape=shape, gen=generators[(op, input_class)]),) + kwargs = {"atol": 0.0, "rtol": 0.0} + cases.append( Case( - name="nodim_3d", - construct={"keepdim": False}, - inputs=((S, S1, S2),), - ), - ], + name=name, + construct={"keepdim": keepdim, "dim": dim}, + inputs=inputs, + **kwargs, + ) + ) + return WebGPUTestSuite( + module_factory=module_cls, + cases=cases, ) @register_op_test("amax") def _amax_suite() -> WebGPUTestSuite: - return _reduce_suite(AmaxModule) + return _reduce_suite(AmaxModule, "amax") @register_op_test("amin") def _amin_suite() -> WebGPUTestSuite: - return _reduce_suite(AminModule) + return _reduce_suite(AminModule, "amin") @register_op_test("flip") @@ -845,6 +884,38 @@ def _view_copy_suite() -> WebGPUTestSuite: return _fn_config_suite(ViewModule, _VIEW_CONFIGS) +def _to_copy_factory(variant: str) -> torch.nn.Module: + return { + "int_to_float": ToCopyIntToFloatModule, + "float_roundtrip": ToCopyFloatToIntToFloatModule, + }[variant]() + + +@register_op_test("to_copy") +def _to_copy_suite() -> WebGPUTestSuite: + cases = [] + for n in (63, 64, 65, 257): + cases.extend( + [ + Case( + name=f"int_to_float_{n}", + construct={"variant": "int_to_float"}, + inputs=(InputSpec(shape=(n,), gen=to_copy_int_input),), + ), + Case( + name=f"float_roundtrip_{n}", + construct={"variant": "float_roundtrip"}, + inputs=(InputSpec(shape=(n,), gen=to_copy_float_input),), + ), + ] + ) + return WebGPUTestSuite( + module_factory=_to_copy_factory, + cases=cases, + golden_dtype="float32", + ) + + @register_op_test("select") def _select_suite() -> WebGPUTestSuite: return _fn_config_suite(SelectModule, _SELECT_CONFIGS) diff --git a/backends/webgpu/test/op_tests/generate_op_tests.py b/backends/webgpu/test/op_tests/generate_op_tests.py index 72f819f94ce..860f0661913 100644 --- a/backends/webgpu/test/op_tests/generate_op_tests.py +++ b/backends/webgpu/test/op_tests/generate_op_tests.py @@ -39,6 +39,8 @@ def _materialize(spec) -> torch.Tensor: shape, gen = spec, "randn" if callable(gen): _t = gen(shape) + if _t.dtype == torch.bool: + return _t return ( _t.to(torch.int32) if not _t.is_floating_point() else _t.to(torch.float32) ) @@ -169,7 +171,10 @@ def generate_case(op: str, suite: WebGPUTestSuite, case, out_dir: str) -> list[d input_entries: list[dict] = [] for i, t in enumerate(inputs): rel = f"{case_id}.in{i}.bin" - if t.dtype == torch.int32: + if t.dtype == torch.bool: + _write_int8(t.to(torch.int8), os.path.join(out_dir, rel)) + in_dtype = "bool" + elif t.dtype == torch.int32: t.detach().cpu().numpy().astype(" +#include #include #include #include @@ -61,7 +62,15 @@ class OpCase : public ::testing::Test { const size_t n = numel(in.shape); std::vector sizes( in.shape.begin(), in.shape.end()); - if (in.dtype == "int32") { + if (in.dtype == "bool") { + auto data = load_int8_bin(in.path, n); + ASSERT_FALSE(data.empty()) << "missing/short input: " << in.path; + std::vector raw(data.begin(), data.end()); + tensors.push_back(make_tensor_ptr( + std::move(sizes), + std::move(raw), + executorch::aten::ScalarType::Bool)); + } else if (in.dtype == "int32") { auto data = load_int32_bin(in.path, n); ASSERT_FALSE(data.empty()) << "missing/short input: " << in.path; tensors.push_back(make_tensor_ptr(std::move(sizes), std::move(data))); @@ -96,10 +105,11 @@ class OpCase : public ::testing::Test { auto golden = load_int8_bin(e_.golden.path, gn); ASSERT_FALSE(golden.empty()) << "missing/short golden: " << e_.golden.path; - const bool* out_p = out_tensor.const_data_ptr(); + ASSERT_EQ(out_tensor.scalar_type(), executorch::aten::ScalarType::Bool); + const uint8_t* out_p = out_tensor.const_data_ptr(); int mism = -1; for (size_t i = 0; i < gn; i++) { - if (static_cast(out_p[i]) != golden[i]) { + if (out_p[i] != static_cast(golden[i])) { mism = static_cast(i); break; } diff --git a/backends/webgpu/test/op_tests/test_generator.py b/backends/webgpu/test/op_tests/test_generator.py index fabf79171c0..65f765812be 100644 --- a/backends/webgpu/test/op_tests/test_generator.py +++ b/backends/webgpu/test/op_tests/test_generator.py @@ -10,7 +10,15 @@ import torch from executorch.backends.webgpu.test.op_tests import generate_op_tests as g -from executorch.backends.webgpu.test.op_tests.test_suite import op_test_registry +from executorch.backends.webgpu.test.op_tests.test_suite import ( + InputSpec, + op_test_registry, +) +from executorch.backends.webgpu.test.ops.test_logical_and import ( + LOGICAL_BINARY_CASES, + logical_binary_gen_a, + logical_binary_gen_b, +) def _add_regular_case(): @@ -112,3 +120,108 @@ def test_manifest_schema_roundtrip(tmp_path): gd = e["golden"] assert {"path", "shape", "dtype", "output_index"} <= set(gd) assert gd["output_index"] == 0 + + +def test_logical_binary_case_contract(): + expected_cases = ( + ("2d", (4, 8)), + ("3d", (2, 3, 8)), + ("sq", (16, 16)), + ("words63", (252,)), + ("words64", (256,)), + ("words65", (260,)), + ) + assert LOGICAL_BINARY_CASES == expected_cases + assert logical_binary_gen_a((8,)).tolist() == [ + -1.0, + -1.0, + 1.0, + 1.0, + -1.0, + -1.0, + 1.0, + 1.0, + ] + assert logical_binary_gen_b((8,)).tolist() == [ + -1.0, + 1.0, + -1.0, + 1.0, + -1.0, + 1.0, + -1.0, + 1.0, + ] + + for op in ("logical_and", "bitwise_and", "logical_or", "bitwise_or"): + suite = op_test_registry[op] + assert tuple((case.name, case.construct["shape"]) for case in suite.cases) == ( + expected_cases + ) + for case in suite.cases: + assert case.required is True + assert case.heavy is False + assert len(case.inputs) == 2 + assert case.inputs[0].gen is logical_binary_gen_a + assert case.inputs[1].gen is logical_binary_gen_b + + +def test_binary_shader_family_case_contract(): + expected = { + "minimum": ( + ("2d", ((37, 41), (37, 41))), + ("3d", ((5, 7, 11), (5, 7, 11))), + ("broadcast_3d_2d", ((2, 3, 8), (3, 1))), + ), + "pow": ( + ("2d", ((37, 41), (37, 41))), + ("3d", ((5, 7, 11), (5, 7, 11))), + ("broadcast_3d_2d", ((2, 3, 8), (3, 1))), + ), + "floor_divide": ( + ("2d", ((37, 41), (37, 41))), + ("3d", ((5, 7, 11), (5, 7, 11))), + ("broadcast_3d_2d", ((2, 3, 8), (3, 1))), + ), + "mul": ( + ("same", ((8, 32), (8, 32))), + ("bcast_lastdim", ((1, 1, 7, 896), (1, 1, 7, 1))), + ("bcast_firstdim", ((4, 4), (1, 4))), + ("bcast_4d_mixed", ((3, 5, 7, 11), (1, 5, 1, 11))), + ("mixedrank", ((4,), (3, 4))), + ), + } + + for op, cases in expected.items(): + suite = op_test_registry[op] + actual = tuple( + ( + case.name, + tuple( + spec.shape if isinstance(spec, InputSpec) else spec + for spec in case.inputs + ), + ) + for case in suite.cases + ) + assert actual == cases + + ranges = { + "minimum": ((-3.0, 3.0), (-2.0, 4.0)), + "pow": ((0.1, 3.0), (-2.0, 3.0)), + "floor_divide": ((-8.0, 8.0), (0.5, 4.0)), + } + for op, expected_ranges in ranges.items(): + case = next( + c for c in op_test_registry[op].cases if c.name == "broadcast_3d_2d" + ) + for spec, (start, end) in zip(case.inputs, expected_ranges): + assert isinstance(spec, InputSpec) and callable(spec.gen) + values = spec.gen(spec.shape).flatten() + assert torch.isclose(values[0], torch.tensor(start)) + assert torch.isclose(values[-1], torch.tensor(end)) + + assert all( + case.golden_fn is g.cases._floor_div_golden + for case in op_test_registry["floor_divide"].cases + ) diff --git a/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py b/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py index 19cb41451e0..b302b3c120d 100644 --- a/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py +++ b/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py @@ -13,15 +13,18 @@ goldens. """ +import math import os import unittest import torch from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner from executorch.exir import to_edge_transform_and_lower +from executorch.exir.backend.utils import get_delegates, get_non_lowered_nodes MAXS = 128 # upper bound for the dynamic seq-len dim (within the 1D dispatch cap) HIDDEN = 64 +_Q4_DECODE_ROUTE_SPEC = "webgpu_record_q4gsw_decode_route" def _rms(x: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6) -> torch.Tensor: @@ -94,6 +97,83 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return torch.sigmoid(x) +class SwiGluModule(torch.nn.Module): + def __init__( + self, + reverse_inner: bool = False, + reverse_outer: bool = False, + extra_gate_consumer: bool = False, + extra_sigmoid_consumer: bool = False, + extra_silu_consumer: bool = False, + expose_intermediates: bool = False, + graph_output: str = "", + separate_inputs: bool = False, + interleaved_projection: bool = False, + qkv_overlap: bool = False, + width: int = 8192, + input_width: int = 64, + group_size: int = 32, + ) -> None: + super().__init__() + from torchao.quantization.granularity import PerGroup + from torchao.quantization.quant_api import IntxWeightOnlyConfig, quantize_ + + def make_q4(seed: int, output_width: int = width): + torch.manual_seed(seed) + linear = torch.nn.Linear(input_width, output_width, bias=False).eval() + quantize_( + linear, + IntxWeightOnlyConfig( + weight_dtype=torch.int4, granularity=PerGroup(group_size) + ), + ) + return linear + + self.gate_proj = make_q4(0) + self.up_proj = make_q4(1) + self.interleaved_proj = make_q4(2) if interleaved_projection else None + self.overlap_q_proj = make_q4(3, 2048) if qkv_overlap else None + self.reverse_inner = reverse_inner + self.reverse_outer = reverse_outer + self.extra_gate_consumer = extra_gate_consumer + self.extra_sigmoid_consumer = extra_sigmoid_consumer + self.extra_silu_consumer = extra_silu_consumer + self.expose_intermediates = expose_intermediates + self.graph_output = graph_output + self.separate_inputs = separate_inputs + self.interleaved_projection = interleaved_projection + self.qkv_overlap = qkv_overlap + self.input_width = input_width + + def forward(self, x: torch.Tensor, up_input: torch.Tensor | None = None): + overlap_q = torch.sigmoid(self.overlap_q_proj(x)) if self.qkv_overlap else None + gate = self.gate_proj(x) + up = self.up_proj(up_input if self.separate_inputs else x) + sigmoid = torch.sigmoid(gate) + silu = sigmoid * gate if self.reverse_inner else gate * sigmoid + interleaved = self.interleaved_proj(x) if self.interleaved_projection else None + output = up * silu if self.reverse_outer else silu * up + if self.extra_gate_consumer: + return output + gate + if self.extra_sigmoid_consumer: + return output + sigmoid + if self.extra_silu_consumer: + return output + silu + if self.expose_intermediates: + return output, gate, sigmoid, silu + if self.graph_output == "gate": + return output, gate + if self.graph_output == "sigmoid": + return output, sigmoid + if self.graph_output == "silu": + return output, silu + if interleaved is not None: + return output + interleaved + if overlap_q is not None: + return overlap_q, output + return output + + class SelectModule(torch.nn.Module): """x.select(0, -1) — negative index resolved live + dynamic output dispatch.""" @@ -108,17 +188,51 @@ def _ramp(shape) -> torch.Tensor: return torch.linspace(-1.0, 1.0, n, dtype=torch.float32).reshape(shape) +def _lower_fully_delegated( + ep, + label: str, + compile_options=None, + *, + expect_q4_decode_route: bool = False, +): + edge = to_edge_transform_and_lower( + ep, + partitioner=[VulkanPartitioner(compile_options=compile_options)], + ) + graph = edge.exported_program().graph_module.graph + delegates = get_delegates(graph) + portable = get_non_lowered_nodes(graph) + if len(delegates) != 1: + raise RuntimeError(f"{label}: expected one delegate, got {len(delegates)}") + if portable: + raise RuntimeError(f"{label}: non-lowered nodes: {portable}") + et = edge.to_executorch() + delegate_ids = [ + delegate.id + for plan in et.executorch_program.execution_plan + for delegate in plan.delegates + ] + if delegate_ids != ["VulkanBackend"]: + raise RuntimeError(f"{label}: serialized delegates: {delegate_ids}") + route_values = [ + spec.value + for plan in et.executorch_program.execution_plan + for delegate in plan.delegates + for spec in delegate.compile_specs + if spec.key == _Q4_DECODE_ROUTE_SPEC + ] + expected_route_values = [b"\x01"] if expect_q4_decode_route else [] + if route_values != expected_route_values: + raise RuntimeError( + f"{label}: {_Q4_DECODE_ROUTE_SPEC}: " + f"expected {expected_route_values}, got {route_values}" + ) + return et + + def _export(model, example_inputs, dynamic_shapes, path: str) -> None: ep = torch.export.export(model, example_inputs, dynamic_shapes=dynamic_shapes) - et = to_edge_transform_and_lower( - ep, partitioner=[VulkanPartitioner()] - ).to_executorch() - found = any( - d.id == "VulkanBackend" - for plan in et.executorch_program.execution_plan - for d in plan.delegates - ) - assert found, f"Expected VulkanBackend delegate in {path}" + et = _lower_fully_delegated(ep, path) with open(path, "wb") as f: f.write(et.buffer) print(f"Exported {path}") @@ -183,16 +297,40 @@ def export_dynamic_shape_cases(out_dir: str) -> None: ) _write_goldens(rmsmul, "dyn_rmsmul", out_dir, [MAXS, 32, 1]) - # 2d) 4-bit quantized linear with a DYNAMIC rows (M) dim — prefill GEMM + # 2d) fp32 linear with a dynamic rows (M) dim. + export_dynamic_fp32_linear_cases(out_dir) + + # 2d.1) 4-bit quantized linear with a DYNAMIC rows (M) dim — prefill GEMM # (register-tiled N=128) + a shmem-GEMM-routed variant (N=2048). _export_dynamic_linear(out_dir) - _export_dynamic_linear(out_dir, n=LIN_SHMEM_N, prefix="dyn_linear_shmem") + _export_dynamic_linear( + out_dir, + n=LIN_SHMEM_N, + prefix="dyn_linear_shmem", + k=LIN_ALT_K, + group=LIN_ALT_GROUP, + ) + _export_dynamic_linear( + out_dir, + prefix="dyn_linear_tiled", + k=LIN_ALT_K, + group=LIN_ALT_GROUP, + ) + _export_static_linear(out_dir, 1, "static_linear_m1") + _export_static_linear(out_dir, 32, "static_linear_m32") + _export_dynamic_bk64_linear_cases(out_dir) # 2e) Fused SDPA with a DYNAMIC seq-len S (prefill, input_pos=0). _export_dynamic_sdpa(out_dir) + _export_dynamic_k16_sdpa_cases(out_dir) + _export_combined_routes(out_dir) + _export_dynamic_qkv_routes(out_dir) + _export_dynamic_sdpa_wide(out_dir) + _export_static_sdpa(out_dir, 1, "static_sdpa_s1") + _export_static_sdpa(out_dir, 16, "static_sdpa_s16") # 2f) 4-bit embedding with a DYNAMIC token count (int64 indices). - _export_dynamic_embedding(out_dir) + export_dynamic_embedding_cases(out_dir) # 2g) Interleaved RoPE with a DYNAMIC seq-len S (two outputs xq/xk). _export_dynamic_rope(out_dir) @@ -207,7 +345,12 @@ def export_dynamic_shape_cases(out_dir: str) -> None: ) _write_goldens(sig, "dyn_sigmoid", out_dir, [MAXS, 32, 1]) - # 2i) select_copy(0, -1) over a DYNAMIC seq-len S (negative live index). + # 2i) Dynamic SwiGLU graph patterns. The Llama-width fixture forces a 2D + # fused dispatch at M=512; compact variants isolate commutative matching and + # graph/intermediate ownership guards around same-input q4 projections. + _export_dynamic_swiglu(out_dir) + + # 2j) select_copy(0, -1) over a DYNAMIC seq-len S (negative live index). _export_dynamic_select(out_dir) # 3) Static rms_norm (no dynamic dim) — regression: must stay byte-identical. @@ -224,36 +367,171 @@ def export_dynamic_shape_cases(out_dir: str) -> None: # Quantized linear: K x N weight, dynamic rows M; input [M, K], output [M, N]. LIN_K = 64 LIN_N = 128 -LIN_SHMEM_N = 2048 # N>=2048 routes linear_q4gsw to the shmem-GEMM path +LIN_ALT_K = 72 # K%16 != 0 disables Steel while K%8 keeps bicol eligible. +LIN_ALT_GROUP = 24 +LIN_SHMEM_N = 2048 LIN_GROUP = 32 LIN_MAXM = 128 +FP32_LINEAR_N = 32 +FP32_LINEAR_MAXM = 128 +FP32_LINEAR_M_VALUES = (FP32_LINEAR_MAXM, 32, 1) + +BK64_K = 2048 +BK64_N = 2048 +BK64_KV_N = 512 +BK64_GATE_N = 8192 +BK64_DOWN_K = 8192 +BK64_GROUP = 64 +BK64_MAXM = 512 +BK64_LIVE_M = (BK64_MAXM, 511, 508, 128, 127, 1) +BK64_OPTIMIZED_M = (BK64_MAXM, 508, 128) +BK64_QKV_LIVE_M = (BK64_MAXM, 511, 508, 128, 127, 16, 2, 1) + + +SWIGLU_MAXM = 512 +SWIGLU_WIDTH = 8192 +SWIGLU_SMALL_WIDTH = 64 +SWIGLU_QKV_OVERLAP_WIDTH = 512 +SWIGLU_K = 64 + + +def _swiglu_inputs(model: SwiGluModule, m: int): + x = _ramp((m, model.input_width)) + if model.separate_inputs: + return x, torch.flip(x, dims=[-1]).contiguous() + return (x,) + + +def _write_swiglu_goldens( + model: SwiGluModule, + prefix: str, + out_dir: str, + m_values, +) -> None: + for m in m_values: + inputs = _swiglu_inputs(model, m) + with torch.no_grad(): + golden = model(*inputs) + base = os.path.join(out_dir, f"{prefix}.S{m}.") + inputs[0].detach().numpy().astype(" None: + inputs = _swiglu_inputs(model, SWIGLU_MAXM) + m_dim = torch.export.Dim("m", min=1, max=SWIGLU_MAXM) + dynamic_shapes = tuple({0: m_dim} for _ in inputs) + ep = torch.export.export( + model.eval(), + inputs, + dynamic_shapes=dynamic_shapes, + ) + et = _lower_fully_delegated(ep, prefix) + with open(os.path.join(out_dir, f"{prefix}.pte"), "wb") as f: + f.write(et.buffer) + print(f"Exported {prefix}.pte") + _write_swiglu_goldens(model, prefix, out_dir, m_values) + + +def _export_dynamic_swiglu(out_dir: str) -> None: + _export_swiglu_case( + out_dir, + "dyn_swiglu", + SwiGluModule(width=SWIGLU_WIDTH), + [SWIGLU_MAXM, 128, 1], + ) + _export_swiglu_case( + out_dir, + "dyn_swiglu_inner_reversed", + SwiGluModule(reverse_inner=True, width=SWIGLU_SMALL_WIDTH), + [128], + ) + for prefix, kwargs in ( + ("dyn_swiglu_extra_sigmoid_consumer", {"extra_sigmoid_consumer": True}), + ("dyn_swiglu_extra_silu_consumer", {"extra_silu_consumer": True}), + ("dyn_swiglu_gate_graph_output", {"graph_output": "gate"}), + ("dyn_swiglu_sigmoid_graph_output", {"graph_output": "sigmoid"}), + ("dyn_swiglu_silu_graph_output", {"graph_output": "silu"}), + ("dyn_swiglu_different_inputs", {"separate_inputs": True}), + ( + "dyn_swiglu_interleaved_q4", + {"interleaved_projection": True}, + ), + ): + _export_swiglu_case( + out_dir, + prefix, + SwiGluModule(width=SWIGLU_SMALL_WIDTH, **kwargs), + [128], + ) + _export_swiglu_case( + out_dir, + "dyn_swiglu_qkv_overlap", + SwiGluModule( + qkv_overlap=True, + width=SWIGLU_QKV_OVERLAP_WIDTH, + input_width=BK64_K, + group_size=BK64_GROUP, + ), + [128], + ) + _export_swiglu_case( + out_dir, + "dyn_swiglu_outer_reversed", + SwiGluModule(reverse_outer=True, width=SWIGLU_SMALL_WIDTH), + [128], + ) + _export_swiglu_case( + out_dir, + "dyn_swiglu_extra_gate_consumer", + SwiGluModule(extra_gate_consumer=True, width=SWIGLU_SMALL_WIDTH), + [128], + ) + _export_swiglu_case( + out_dir, + "dyn_swiglu_graph_outputs", + SwiGluModule(expose_intermediates=True, width=SWIGLU_SMALL_WIDTH), + [128], + ) + def _export_dynamic_linear( - out_dir: str, n: int = LIN_N, prefix: str = "dyn_linear" + out_dir: str, + n: int = LIN_N, + prefix: str = "dyn_linear", + k: int = LIN_K, + group: int = LIN_GROUP, ) -> None: - from executorch.backends.webgpu.test.ops.quantized_linear.test_quantized_linear import ( + from executorch.backends.webgpu.test.ops.test_quantized_linear import ( _fp64_golden, _make_quantized_model, ) - model = _make_quantized_model(LIN_K, n, LIN_GROUP) - x = _ramp((LIN_MAXM, LIN_K)) + model = _make_quantized_model(k, n, group) + x = _ramp((LIN_MAXM, k)) m_dim = torch.export.Dim("m", min=1, max=LIN_MAXM) ep = torch.export.export(model, (x,), dynamic_shapes=({0: m_dim},)) - et = to_edge_transform_and_lower( - ep, partitioner=[VulkanPartitioner()] - ).to_executorch() - assert any( - d.id == "VulkanBackend" - for plan in et.executorch_program.execution_plan - for d in plan.delegates - ), "linear_q4gsw not delegated" + et = _lower_fully_delegated(ep, prefix) with open(os.path.join(out_dir, f"{prefix}.pte"), "wb") as f: f.write(et.buffer) print(f"Exported {prefix}.pte") for m in [LIN_MAXM, 32, 1]: - xm = _ramp((m, LIN_K)) + xm = _ramp((m, k)) g = _fp64_golden(model, xm).astype(" None: + from executorch.backends.webgpu.test.ops.test_linear_fp32 import make_linear + + model = make_linear(k, FP32_LINEAR_N, bias=bias).eval() + x = _ramp((FP32_LINEAR_MAXM, k)) + m_dim = torch.export.Dim("m", min=1, max=FP32_LINEAR_MAXM) + ep = torch.export.export(model, (x,), dynamic_shapes=({0: m_dim},)) + et = _lower_fully_delegated(ep, prefix) + with open(os.path.join(out_dir, f"{prefix}.pte"), "wb") as f: + f.write(et.buffer) + + weight = model.fc.weight.detach().double() + bias_value = model.fc.bias.detach().double() if model.fc.bias is not None else None + for m in FP32_LINEAR_M_VALUES: + xm = _ramp((m, k)) + golden = torch.nn.functional.linear(xm.double(), weight, bias_value) + base = os.path.join(out_dir, f"{prefix}.S{m}") + xm.detach().numpy().astype(" None: + os.makedirs(out_dir, exist_ok=True) + for route, k in (("vec4", 64), ("tiled", 63)): + for bias in (True, False): + suffix = "bias" if bias else "no_bias" + _export_dynamic_fp32_linear_case( + out_dir, + k=k, + bias=bias, + prefix=f"dyn_linear_fp32_{route}_{suffix}", + ) + + +def _make_bk64_model( + *, + k: int = BK64_K, + n: int = BK64_N, + group: int = BK64_GROUP, + bias: bool = False, + seed: int = 11, +) -> torch.nn.Module: + from torchao.quantization.granularity import PerGroup + from torchao.quantization.quant_api import IntxWeightOnlyConfig, quantize_ + + torch.manual_seed(seed) + model = torch.nn.Linear(k, n, bias=bias).eval() + if model.bias is not None: + with torch.no_grad(): + model.bias.copy_(torch.linspace(-0.25, 0.25, n)) + quantize_( + model, + IntxWeightOnlyConfig(weight_dtype=torch.int4, granularity=PerGroup(group)), + ) + return model + + +def _bk64_golden(model: torch.nn.Module, x: torch.Tensor) -> torch.Tensor: + golden = x.double() @ model.weight.dequantize().double().t() + if model.bias is not None: + golden = golden + model.bias.double() + return golden.to(torch.float32) + + +def _bk64_input(m: int, k: int) -> torch.Tensor: + flat = torch.arange(m * k, dtype=torch.int64) + hashed = (flat * 37 + torch.div(flat, 16, rounding_mode="floor") * 53) % 257 + return ((hashed.to(torch.float32) - 128.0) / 128.0).reshape(m, k) + + +class Bk64ShapeAwareLinear(torch.nn.Module): + def __init__(self, projection: torch.nn.Module, output_width: int) -> None: + super().__init__() + self.projection = projection + self.output_width = output_width + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.projection(x).reshape(1, x.shape[0], self.output_width) + + +class Bk64Qkv(torch.nn.Module): + def __init__( + self, + *, + widths=(BK64_N, BK64_KV_N, BK64_KV_N), + group: int = BK64_GROUP, + bias: bool = False, + separate_v_input: bool = False, + ) -> None: + super().__init__() + self.q = _make_bk64_model(n=widths[0], group=group, bias=bias, seed=21) + self.k = _make_bk64_model(n=widths[1], group=group, bias=bias, seed=22) + self.v = _make_bk64_model(n=widths[2], group=group, bias=bias, seed=23) + self.widths = widths + self.separate_v_input = separate_v_input + + def forward( + self, + x: torch.Tensor, + v_input: torch.Tensor | None = None, + ): + q = self.q(x).reshape(1, x.shape[0], self.widths[0]) + k = self.k(x).reshape(1, x.shape[0], self.widths[1]) + v_source = v_input if self.separate_v_input else x + v = self.v(v_source).reshape(1, v_source.shape[0], self.widths[2]) + return q, k, v + + +def _export_bk64_program( + model: torch.nn.Module, + x: torch.Tensor, + n: int, + prefix: str, +): + export_model = Bk64ShapeAwareLinear(model, n).eval() + m_dim = torch.export.Dim("m", min=1, max=BK64_MAXM) + ep = torch.export.export(export_model, (x,), dynamic_shapes=({0: m_dim},)) + if not any(node.target == torch.ops.aten.sym_size.int for node in ep.graph.nodes): + raise RuntimeError(f"{prefix}: dynamic q4 fixture lost aten.sym_size.int") + return ep + + +def _export_dynamic_bk64_linear_case( + out_dir: str, + prefix: str, + *, + k: int = BK64_K, + n: int = BK64_N, + group: int = BK64_GROUP, + bias: bool = False, + live_m=BK64_LIVE_M, +) -> None: + model = _make_bk64_model(k=k, n=n, group=group, bias=bias) + x = _bk64_input(BK64_MAXM, k) + ep = _export_bk64_program(model, x, n, prefix) + et = _lower_fully_delegated(ep, prefix) + with open(os.path.join(out_dir, f"{prefix}.pte"), "wb") as f: + f.write(et.buffer) + print(f"Exported {prefix}.pte") + for m in live_m: + xm = _bk64_input(m, k) + golden = _bk64_golden(model, xm) + xm.detach().numpy().astype(" None: + os.makedirs(out_dir, exist_ok=True) + _export_dynamic_bk64_linear_case(out_dir, "dyn_linear_bk64") + _export_dynamic_bk64_linear_case( + out_dir, + "dyn_linear_bk64_gate", + n=BK64_GATE_N, + live_m=BK64_OPTIMIZED_M, + ) + _export_dynamic_bk64_linear_case( + out_dir, + "dyn_linear_bk64_down", + k=BK64_DOWN_K, + live_m=BK64_OPTIMIZED_M, + ) + _export_dynamic_bk64_linear_case( + out_dir, + "dyn_linear_bk64_group32", + group=32, + live_m=[128], + ) + _export_dynamic_bk64_linear_case( + out_dir, + "dyn_linear_bk64_bias", + bias=True, + live_m=[128], + ) + _export_dynamic_bk64_linear_case( + out_dir, + "dyn_linear_bk64_kv_shape", + n=BK64_KV_N, + live_m=[128], + ) + _export_dynamic_bk64_qkv_cases(out_dir) + + +def _export_dynamic_bk64_qkv_case( + out_dir: str, + prefix: str, + model: Bk64Qkv, + live_m, +) -> None: + inputs = (_bk64_input(BK64_MAXM, BK64_K),) + if model.separate_v_input: + inputs += (torch.flip(inputs[0], dims=[-1]).contiguous(),) + m_dim = torch.export.Dim("m", min=1, max=BK64_MAXM) + dynamic_shapes = tuple({0: m_dim} for _ in inputs) + ep = torch.export.export(model.eval(), inputs, dynamic_shapes=dynamic_shapes) + if not any(node.target == torch.ops.aten.sym_size.int for node in ep.graph.nodes): + raise RuntimeError(f"{prefix}: dynamic QKV fixture lost aten.sym_size.int") + et = _lower_fully_delegated(ep, prefix) + with open(os.path.join(out_dir, f"{prefix}.pte"), "wb") as f: + f.write(et.buffer) + print(f"Exported {prefix}.pte") + for m in live_m: + x = _bk64_input(m, BK64_K) + case_inputs = (x,) + if model.separate_v_input: + case_inputs += (torch.flip(x, dims=[-1]).contiguous(),) + outputs = ( + _bk64_golden(model.q, case_inputs[0]), + _bk64_golden(model.k, case_inputs[0]), + _bk64_golden(model.v, case_inputs[-1]), + ) + base = os.path.join(out_dir, f"{prefix}.S{m}.") + case_inputs[0].detach().numpy().astype(" None: + _export_dynamic_bk64_qkv_case( + out_dir, + "dyn_qkv_bk64", + Bk64Qkv(), + BK64_QKV_LIVE_M, + ) + for prefix, model in ( + ("dyn_qkv_bk64_group32", Bk64Qkv(group=32)), + ("dyn_qkv_bk64_bias", Bk64Qkv(bias=True)), + ( + "dyn_qkv_bk64_wrong_width", + Bk64Qkv(widths=(BK64_N, BK64_N, BK64_KV_N)), + ), + ( + "dyn_qkv_bk64_different_input", + Bk64Qkv(separate_v_input=True), + ), + ): + _export_dynamic_bk64_qkv_case(out_dir, prefix, model, [128]) + + +def _export_static_linear(out_dir: str, m: int, prefix: str) -> None: + from executorch.backends.webgpu.test.ops.test_quantized_linear import ( + _fp64_golden, + _make_quantized_model, + ) + + model = _make_quantized_model(LIN_K, LIN_N, LIN_GROUP) + x = _ramp((m, LIN_K)) + ep = torch.export.export(model, (x,)) + et = _lower_fully_delegated(ep, prefix) + with open(os.path.join(out_dir, f"{prefix}.pte"), "wb") as f: + f.write(et.buffer) + x.detach().numpy().astype(" None: - from executorch.backends.webgpu.test.ops.sdpa.test_sdpa import ( + from executorch.backends.webgpu.test.ops.test_sdpa import ( _det_inputs, _golden, SdpaConfig, @@ -286,14 +852,7 @@ def cfg(s: int) -> "SdpaConfig": s_dim = torch.export.Dim("s", min=1, max=SD_MAXS) ds = ({1: s_dim}, {1: s_dim}, {1: s_dim}, None, None) ep = torch.export.export(model, (q, k, v, kc, vc), dynamic_shapes=ds) - et = to_edge_transform_and_lower( - ep, partitioner=[VulkanPartitioner()] - ).to_executorch() - assert any( - d.id == "VulkanBackend" - for plan in et.executorch_program.execution_plan - for d in plan.delegates - ), "sdpa not delegated" + et = _lower_fully_delegated(ep, "sdpa_dyn") with open(os.path.join(out_dir, "sdpa_dyn.pte"), "wb") as f: f.write(et.buffer) print("Exported sdpa_dyn.pte") @@ -315,6 +874,478 @@ def cfg(s: int) -> "SdpaConfig": print(f" golden sdpa_dyn S={s} (golden shape {tuple(g.shape)})") +def _write_f32_tensors(base, tensors) -> None: + for name, tensor in tensors: + tensor.detach().numpy().astype(" None: + if not initialized_cache: + k_cache.zero_() + v_cache.zero_() + + +def _export_dynamic_k16_sdpa_case( + out_dir: str, + prefix: str, + hq: int, + hkv: int, + live_s, + d: int = K16_D, + scale: float | None = None, + cache_hkv: int | None = None, + expect_runtime_reject: bool = False, + denom: float = 16.0, + kv_f16_golden: bool = False, + initialized_cache: bool = False, +) -> None: + from executorch.backends.webgpu.test.ops.test_sdpa import ( + _det_inputs, + _golden, + _round_kv_for_storage, + SdpaConfig, + ) + + def cfg(s: int) -> "SdpaConfig": + return SdpaConfig( + prefix, + hq, + hkv, + d, + s, + K16_CMAX, + K16_INPUT_POS, + denom, + kv_f16=kv_f16_golden, + ) + + q, k, v, kc, vc = _det_inputs(cfg(K16_MAXS)) + if cache_hkv is not None: + kc = torch.zeros(1, K16_CMAX, cache_hkv, d) + vc = torch.zeros_like(kc) + _zero_uninitialized_cache(initialized_cache, kc, vc) + + class K16SdpaModule(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.register_buffer("k_cache", kc) + self.register_buffer("v_cache", vc) + + def forward(self, q, k, v, pos_control): + input_pos = pos_control.shape[1] + return torch.ops.llama.sdpa_with_kv_cache( + q, + k, + v, + self.k_cache, + self.v_cache, + input_pos, + q.shape[1], + None, + 0.0, + True, + scale, + ) + + model = K16SdpaModule().eval() + inputs = (q, k, v, torch.zeros(1, K16_POS_CONTROL)) + s_dim = torch.export.Dim(f"{prefix}_s", min=1, max=K16_MAXS) + pos_dim = torch.export.Dim(f"{prefix}_pos_control", min=1, max=K16_POS_CONTROL) + ep = torch.export.export( + model, + inputs, + dynamic_shapes=({1: s_dim}, {1: s_dim}, {1: s_dim}, {1: pos_dim}), + ) + sym_sizes = [ + node for node in ep.graph.nodes if node.target == torch.ops.aten.sym_size.int + ] + if len(sym_sizes) < 2: + raise RuntimeError(f"{prefix}: dynamic K16 fixture lost S/position symbols") + et = _lower_fully_delegated(ep, prefix) + with open(os.path.join(out_dir, f"{prefix}.pte"), "wb") as f: + f.write(et.buffer) + print(f"Exported {prefix}.pte") + if expect_runtime_reject: + return + + def reference(live_cfg, q, k, v, kc, vc): + if scale is None: + return _golden(live_cfg, q, k, v, kc, vc) + k, v, kc, vc = _round_kv_for_storage(live_cfg, k, v, kc, vc) + runtime_scale = float(torch.tensor(scale, dtype=torch.float32).item()) + context_len = live_cfg.s + live_cfg.input_pos + g = hq // hkv + k_full = torch.cat((kc[0, : live_cfg.input_pos].double(), k[0].double()), dim=0) + v_full = torch.cat((vc[0, : live_cfg.input_pos].double(), v[0].double()), dim=0) + q_heads = q[0].double().transpose(0, 1) + k_heads = k_full.repeat_interleave(g, dim=1).transpose(0, 1) + v_heads = v_full.repeat_interleave(g, dim=1).transpose(0, 1) + mask = torch.full((live_cfg.s, context_len), float("-inf"), dtype=torch.float64) + for token in range(live_cfg.s): + mask[token, : live_cfg.input_pos + token + 1] = 0.0 + golden = torch.nn.functional.scaled_dot_product_attention( + q_heads, k_heads, v_heads, attn_mask=mask, scale=runtime_scale + ) + return golden.transpose(0, 1).reshape(1, live_cfg.s, hq, d).float().contiguous() + + if initialized_cache: + initial_cfg = cfg(17) + initial_q, initial_k, initial_v, initial_kc, initial_vc = _det_inputs( + initial_cfg + ) + initial_golden = reference( + initial_cfg, + initial_q, + initial_k, + initial_v, + initial_kc, + initial_vc, + ) + initial_base = os.path.join(out_dir, f"{prefix}.initial.") + _write_f32_tensors( + initial_base, + ( + ("q", initial_q), + ("k", initial_k), + ("v", initial_v), + ("control", torch.zeros(1, K16_POS_CONTROL)), + ("golden", initial_golden), + ), + ) + + prime_cfg = SdpaConfig( + prefix, + hq, + hkv, + d, + K16_PRIME_S, + K16_CMAX, + K16_PRIME_POS, + denom, + kv_f16=kv_f16_golden, + ) + prime_q, prime_k, prime_v, prime_kc, prime_vc = _det_inputs(prime_cfg) + _zero_uninitialized_cache(initialized_cache, prime_kc, prime_vc) + prime_golden = reference(prime_cfg, prime_q, prime_k, prime_v, prime_kc, prime_vc) + prime_base = os.path.join(out_dir, f"{prefix}.prime.") + _write_f32_tensors( + prime_base, + ( + ("q", prime_q), + ("k", prime_k), + ("v", prime_v), + ("control", torch.zeros(1, 1)), + ("golden", prime_golden), + ), + ) + + for s in live_s: + live_cfg = cfg(s) + q, k, v, kc, vc = _det_inputs(live_cfg) + _zero_uninitialized_cache(initialized_cache, kc, vc) + kc[0, K16_PRIME_POS:K16_INPUT_POS] = prime_k[0] + vc[0, K16_PRIME_POS:K16_INPUT_POS] = prime_v[0] + golden = reference(live_cfg, q, k, v, kc, vc) + base = os.path.join(out_dir, f"{prefix}.S{s}.") + _write_f32_tensors( + base, + ( + ("q", q), + ("k", k), + ("v", v), + ("control", torch.zeros(1, K16_POS_CONTROL)), + ("kc", kc), + ("vc", vc), + ("golden", golden), + ), + ) + print(f" golden {prefix} S={s}") + + +def _export_dynamic_k16_sdpa_cases(out_dir: str) -> None: + _export_dynamic_k16_sdpa_case( + out_dir, + "sdpa_k16_llama", + K16_HQ, + K16_HKV, + K16_LIVE_S, + ) + _export_dynamic_k16_sdpa_case( + out_dir, + "sdpa_k16_qwen3", + QWEN3_HQ, + QWEN3_HKV, + QWEN3_LIVE_S, + d=QWEN3_D, + denom=10.0, + kv_f16_golden=True, + initialized_cache=True, + ) + _export_dynamic_k16_sdpa_case( + out_dir, + "sdpa_k16_qwen3_near_scale", + QWEN3_HQ, + QWEN3_HKV, + (128, 1), + d=QWEN3_D, + scale=1.0 / math.sqrt(float(QWEN3_D)) + 5e-7, + denom=10.0, + kv_f16_golden=True, + ) + _export_dynamic_k16_sdpa_case( + out_dir, + "sdpa_k16_wrong_geometry", + 14, + 2, + (128, 1), + ) + _export_dynamic_k16_sdpa_case( + out_dir, + "sdpa_k16_wrong_d", + K16_HQ, + K16_HKV, + (128, 1), + d=128, + scale=0.125, + ) + _export_dynamic_k16_sdpa_case( + out_dir, + "sdpa_k16_wrong_scale", + K16_HQ, + K16_HKV, + (128, 1), + scale=0.25, + ) + _export_dynamic_k16_sdpa_case( + out_dir, + "sdpa_k16_bad_cache_heads", + K16_HQ, + K16_HKV, + (), + cache_hkv=K16_HKV - 1, + expect_runtime_reject=True, + ) + + +def _export_combined_routes(out_dir: str) -> None: + from executorch.backends.webgpu.test.ops.test_quantized_linear import ( + _make_quantized_model, + ) + from executorch.backends.webgpu.test.ops.test_sdpa import ( + _det_inputs, + _golden, + SdpaConfig, + ) + from executorch.extension.llm.custom_ops import custom_ops # noqa: F401 + + class CombinedRoutes(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear = _make_quantized_model(LIN_K, SD_HQ * SD_D, LIN_GROUP) + + def forward(self, x, q, k, v, k_cache, v_cache): + projected = self.linear(x).reshape(1, x.shape[0], SD_HQ, SD_D) + return torch.ops.llama.sdpa_with_kv_cache( + q + projected, + k, + v, + k_cache, + v_cache, + 0, + q.shape[1], + None, + 0.0, + True, + None, + ) + + model = CombinedRoutes().eval() + cfg = SdpaConfig("combined", SD_HQ, SD_HKV, SD_D, SD_MAXS, SD_CMAX, 0) + q, k, v, kc, vc = _det_inputs(cfg) + x = _ramp((SD_MAXS, LIN_K)) + s_dim = torch.export.Dim("s", min=1, max=SD_MAXS) + ep = torch.export.export( + model, + (x, q, k, v, kc, vc), + dynamic_shapes=( + {0: s_dim}, + {1: s_dim}, + {1: s_dim}, + {1: s_dim}, + None, + None, + ), + ) + et = _lower_fully_delegated(ep, "combined_routes") + with open(os.path.join(out_dir, "combined_routes.pte"), "wb") as f: + f.write(et.buffer) + print("Exported combined_routes.pte") + + for s in [SD_MAXS, 16, 1]: + live_cfg = SdpaConfig("combined", SD_HQ, SD_HKV, SD_D, s, SD_CMAX, 0) + q, k, v, kc, vc = _det_inputs(live_cfg) + x = _ramp((s, LIN_K)) + with torch.no_grad(): + projected = model.linear(x).reshape(1, s, SD_HQ, SD_D) + golden = _golden(live_cfg, q + projected, k, v, kc, vc) + base = os.path.join(out_dir, f"combined_routes.S{s}.") + for name, tensor in ( + ("x", x), + ("q", q), + ("k", k), + ("v", v), + ("kc", kc), + ("vc", vc), + ("golden", golden), + ): + tensor.detach().numpy().astype(" None: + from executorch.backends.webgpu.test.ops.test_quantized_linear import ( + _make_quantized_model, + ) + + class QkvRoutes(torch.nn.Module): + def __init__(self, k: int = LIN_K, group: int = LIN_GROUP): + super().__init__() + self.q = _make_quantized_model(k, QKV_NQ, group, seed=0) + self.k = _make_quantized_model(k, QKV_NK, group, seed=1) + self.v = _make_quantized_model(k, QKV_NV, group, seed=2) + + def forward(self, x): + # Keep the linears internal so graph-output copies do not capture + # the buffers that the runtime QKV pass later replaces. + return ( + torch.sigmoid(self.q(x)), + torch.sigmoid(self.k(x)), + torch.sigmoid(self.v(x)), + ) + + model = QkvRoutes().eval() + x = _ramp((QKV_MAXM, LIN_K)) + m_dim = torch.export.Dim("m", min=1, max=QKV_MAXM) + ep = torch.export.export(model, (x,), dynamic_shapes=({0: m_dim},)) + et = _lower_fully_delegated( + ep, + "qkv_routes", + compile_options={_Q4_DECODE_ROUTE_SPEC: True}, + expect_q4_decode_route=True, + ) + with open(os.path.join(out_dir, "qkv_routes.pte"), "wb") as f: + f.write(et.buffer) + print("Exported qkv_routes.pte") + + for m in [QKV_MAXM, 1]: + live_x = _ramp((m, LIN_K)) + with torch.no_grad(): + q, k, v = model(live_x) + base = os.path.join(out_dir, f"qkv_routes.S{m}.") + for name, tensor in (("input", live_x), ("q", q), ("k", k), ("v", v)): + tensor.detach().numpy().astype(" None: + from executorch.backends.webgpu.test.ops.test_sdpa import ( + _det_inputs, + _golden, + SdpaConfig, + SdpaModule, + ) + + hq, hkv, d, cmax, max_s = 8, 2, 132, 16, 16 + + def cfg(s): + return SdpaConfig("wide", hq, hkv, d, s, cmax, 0) + + model = SdpaModule(0) + q, k, v, kc, vc = _det_inputs(cfg(max_s)) + s_dim = torch.export.Dim("s", min=1, max=max_s) + ep = torch.export.export( + model, + (q, k, v, kc, vc), + dynamic_shapes=({1: s_dim}, {1: s_dim}, {1: s_dim}, None, None), + ) + et = _lower_fully_delegated(ep, "sdpa_wide") + with open(os.path.join(out_dir, "sdpa_wide.pte"), "wb") as f: + f.write(et.buffer) + for s in [max_s, 1]: + live = cfg(s) + q, k, v, kc, vc = _det_inputs(live) + golden = _golden(live, q, k, v, kc, vc) + for name, tensor in ( + ("q", q), + ("k", k), + ("v", v), + ("kc", kc), + ("vc", vc), + ("golden", golden), + ): + tensor.detach().numpy().astype(" None: + from executorch.backends.webgpu.test.ops.test_sdpa import ( + _det_inputs, + _golden, + SdpaConfig, + SdpaModule, + ) + + cfg = SdpaConfig(prefix, SD_HQ, SD_HKV, SD_D, s, SD_CMAX, 0) + inputs = _det_inputs(cfg) + ep = torch.export.export(SdpaModule(0), inputs) + et = _lower_fully_delegated(ep, prefix) + with open(os.path.join(out_dir, f"{prefix}.pte"), "wb") as f: + f.write(et.buffer) + golden = _golden(cfg, *inputs) + for name, tensor in zip(("q", "k", "v", "kc", "vc"), inputs): + tensor.detach().numpy().astype(" [N, EMBED] fp32. EMB_VOCAB = 64 EMB_DIM = 64 @@ -322,52 +1353,64 @@ def cfg(s: int) -> "SdpaConfig": EMB_MAXN = 16 -class _LinearPackedEmbedding(torch.nn.Module): - def __init__(self) -> None: +def _write_embedding_goldens( + out_dir: str, + prefix: str, + weight: torch.Tensor, + scales: torch.Tensor, + group_size: int, + is_linear: bool, +) -> None: + for n in [EMB_MAXN, 8, 1]: + idx = (torch.arange(n, dtype=torch.long) * 7) % EMB_VOCAB + golden = torch.ops.et_vk.embedding_q4gsw.default( + weight, scales, group_size, idx, is_linear + ) + idx.detach().numpy().astype(" None: super().__init__() packed = torch.arange(EMB_VOCAB * (EMB_DIM // 2), dtype=torch.int64).reshape( EMB_VOCAB, EMB_DIM // 2 ) self.register_buffer("weight", (packed % 256).to(torch.uint8)) self.register_buffer("scales", torch.ones(EMB_VOCAB, EMB_DIM // EMB_GROUP)) + self.is_linear_weight = is_linear_weight def forward(self, indices: torch.Tensor) -> torch.Tensor: return torch.ops.et_vk.embedding_q4gsw.default( - self.weight, self.scales, EMB_GROUP, indices, True + self.weight, + self.scales, + EMB_GROUP, + indices, + self.is_linear_weight, ) -def _write_embedding_goldens( - out_dir: str, - prefix: str, - weight: torch.Tensor, - scales: torch.Tensor, - group_size: int, - is_linear: bool, +def _write_packed_embedding_goldens( + out_dir: str, prefix: str, model: _PackedEmbedding ) -> None: for n in [EMB_MAXN, 8, 1]: idx = (torch.arange(n, dtype=torch.long) * 7) % EMB_VOCAB - golden = torch.ops.et_vk.embedding_q4gsw.default( - weight, scales, group_size, idx, is_linear - ) - if is_linear: - nonlinear_golden = torch.ops.et_vk.embedding_q4gsw.default( - weight, scales, group_size, idx, False - ) - if torch.equal(golden, nonlinear_golden): - raise RuntimeError( - "emb_dyn_linear fixture does not distinguish nibble packing" - ) + golden = model(idx) idx.detach().numpy().astype(" None: +def export_dynamic_embedding_cases(out_dir: str) -> None: + os.makedirs(out_dir, exist_ok=True) from executorch.backends.webgpu.test.ops.test_embedding_q4gsw import ( _make_quantized_model, _quant_params, @@ -379,35 +1422,29 @@ def _export_dynamic_embedding(out_dir: str) -> None: idx_max = torch.arange(EMB_MAXN, dtype=torch.long) n_dim = torch.export.Dim("n", min=1, max=EMB_MAXN) ep = torch.export.export(qm, (idx_max,), dynamic_shapes=({0: n_dim},)) - et = to_edge_transform_and_lower( - ep, partitioner=[VulkanPartitioner()] - ).to_executorch() - assert any( - d.id == "VulkanBackend" - for plan in et.executorch_program.execution_plan - for d in plan.delegates - ), "embedding_q4gsw not delegated" + et = _lower_fully_delegated(ep, "emb_dyn") with open(os.path.join(out_dir, "emb_dyn.pte"), "wb") as f: f.write(et.buffer) print("Exported emb_dyn.pte") weight, scales, group_size = _quant_params(qm) _write_embedding_goldens(out_dir, "emb_dyn", weight, scales, group_size, False) - linear_model = _LinearPackedEmbedding().eval() - _export( - linear_model, - (idx_max,), - ({0: n_dim},), - os.path.join(out_dir, "emb_dyn_linear.pte"), - ) - _write_embedding_goldens( - out_dir, - "emb_dyn_linear", - linear_model.weight, - linear_model.scales, - EMB_GROUP, - True, - ) + packed_models = { + "emb_dyn_linear": _PackedEmbedding(True).eval(), + "emb_dyn_nonlinear": _PackedEmbedding(False).eval(), + } + linear_golden = packed_models["emb_dyn_linear"](idx_max) + nonlinear_golden = packed_models["emb_dyn_nonlinear"](idx_max) + if torch.equal(linear_golden, nonlinear_golden): + raise RuntimeError("embedding layout fixtures must distinguish nibble order") + for prefix, model in packed_models.items(): + _export( + model, + (idx_max,), + ({0: n_dim},), + os.path.join(out_dir, f"{prefix}.pte"), + ) + _write_packed_embedding_goldens(out_dir, prefix, model) # Dynamic RoPE: xq/xk + freqs all share a dynamic seq-len S. @@ -418,11 +1455,7 @@ def _export_dynamic_embedding(out_dir: str) -> None: def _export_dynamic_rope(out_dir: str) -> None: - from executorch.backends.webgpu.test.ops.rope.test_rope import ( - _golden, - _inputs, - Shape, - ) + from executorch.backends.webgpu.test.ops.test_rope import _golden, _inputs, Shape from executorch.examples.models.llama.rope import RotaryEmbedding xq, xk, fc, fs = _inputs(Shape("dyn", 1, ROPE_MAXS, ROPE_NH, ROPE_NKV, ROPE_HD)) @@ -431,14 +1464,7 @@ def _export_dynamic_rope(out_dir: str) -> None: ep = torch.export.export( RotaryEmbedding().eval(), (xq, xk, fc, fs), dynamic_shapes=ds ) - et = to_edge_transform_and_lower( - ep, partitioner=[VulkanPartitioner()] - ).to_executorch() - assert any( - d.id == "VulkanBackend" - for plan in et.executorch_program.execution_plan - for d in plan.delegates - ), "apply_rotary_emb not delegated" + et = _lower_fully_delegated(ep, "rope_dyn") with open(os.path.join(out_dir, "rope_dyn.pte"), "wb") as f: f.write(et.buffer) print("Exported rope_dyn.pte") @@ -471,14 +1497,7 @@ def _export_dynamic_select(out_dir: str) -> None: (_ramp((SEL_LEAD, 1, MAXS, HIDDEN)),), dynamic_shapes=({2: s_dim},), ) - et = to_edge_transform_and_lower( - ep, partitioner=[VulkanPartitioner()] - ).to_executorch() - assert any( - d.id == "VulkanBackend" - for plan in et.executorch_program.execution_plan - for d in plan.delegates - ), "select_copy not delegated" + et = _lower_fully_delegated(ep, "dyn_select") with open(os.path.join(out_dir, "dyn_select.pte"), "wb") as f: f.write(et.buffer) print("Exported dyn_select.pte") @@ -496,6 +1515,19 @@ def _export_dynamic_select(out_dir: str) -> None: class TestDynamicShapeExport(unittest.TestCase): + def test_q4_route_compile_specs(self) -> None: + import tempfile + + model = RmsNormModule(HIDDEN).eval() + x = _ramp((1, 1, MAXS, HIDDEN)) + ep = torch.export.export(model, (x,)) + _lower_fully_delegated(ep, "q4_route_negative_control") + + with tempfile.TemporaryDirectory() as d: + _export_dynamic_qkv_routes(d) + self.assertTrue(os.path.exists(os.path.join(d, "qkv_routes.pte"))) + self.assertTrue(os.path.exists(os.path.join(d, "qkv_bk64_routes.pte"))) + def test_export_dynamic_rms(self) -> None: import tempfile @@ -503,6 +1535,84 @@ def test_export_dynamic_rms(self) -> None: export_dynamic_shape_cases(d) self.assertTrue(os.path.exists(os.path.join(d, "dyn_rms.pte"))) self.assertTrue(os.path.exists(os.path.join(d, "dyn_rms.S1.golden.bin"))) + expected = [ + "dyn_linear_bk64.pte", + "dyn_linear_bk64.S512.input.bin", + "dyn_linear_bk64.S512.golden.bin", + "dyn_linear_bk64.S511.input.bin", + "dyn_linear_bk64.S511.golden.bin", + "dyn_linear_bk64.S508.golden.bin", + "dyn_linear_bk64.S128.golden.bin", + "dyn_linear_bk64.S127.golden.bin", + "dyn_linear_bk64.S1.golden.bin", + "dyn_linear_bk64_gate.pte", + "dyn_linear_bk64_gate.S512.input.bin", + "dyn_linear_bk64_gate.S512.golden.bin", + "dyn_linear_bk64_gate.S508.golden.bin", + "dyn_linear_bk64_gate.S128.golden.bin", + "dyn_linear_bk64_down.pte", + "dyn_linear_bk64_down.S512.input.bin", + "dyn_linear_bk64_down.S512.golden.bin", + "dyn_linear_bk64_down.S508.golden.bin", + "dyn_linear_bk64_down.S128.golden.bin", + "dyn_linear_bk64_group32.pte", + "dyn_linear_bk64_bias.pte", + "dyn_linear_bk64_kv_shape.pte", + "dyn_qkv_bk64.pte", + "dyn_qkv_bk64.S512.input.bin", + "dyn_qkv_bk64.S512.q.bin", + "dyn_qkv_bk64.S512.k.bin", + "dyn_qkv_bk64.S512.v.bin", + "dyn_qkv_bk64.S511.input.bin", + "dyn_qkv_bk64.S511.q.bin", + "dyn_qkv_bk64.S511.k.bin", + "dyn_qkv_bk64.S511.v.bin", + "dyn_qkv_bk64.S508.q.bin", + "dyn_qkv_bk64.S508.k.bin", + "dyn_qkv_bk64.S508.v.bin", + "dyn_qkv_bk64.S128.q.bin", + "dyn_qkv_bk64.S127.q.bin", + "dyn_qkv_bk64.S16.q.bin", + "dyn_qkv_bk64.S2.q.bin", + "dyn_qkv_bk64.S1.q.bin", + "dyn_qkv_bk64_group32.pte", + "dyn_qkv_bk64_bias.pte", + "dyn_qkv_bk64_wrong_width.pte", + "dyn_qkv_bk64_different_input.pte", + "sdpa_k16_bad_cache_heads.pte", + "sdpa_k16_qwen3.pte", + "sdpa_k16_qwen3_near_scale.pte", + ] + for name in expected: + with self.subTest(artifact=name): + self.assertGreater(os.path.getsize(os.path.join(d, name)), 0) + for prefix, live_s in ( + ("sdpa_k16_llama", K16_LIVE_S), + ("sdpa_k16_qwen3", QWEN3_LIVE_S), + ("sdpa_k16_qwen3_near_scale", (128, 1)), + ("sdpa_k16_wrong_geometry", (128, 1)), + ("sdpa_k16_wrong_d", (128, 1)), + ("sdpa_k16_wrong_scale", (128, 1)), + ): + with self.subTest(artifact=f"{prefix}.pte"): + self.assertGreater( + os.path.getsize(os.path.join(d, f"{prefix}.pte")), 0 + ) + for kind in ("q", "k", "v", "control", "golden"): + name = f"{prefix}.prime.{kind}.bin" + with self.subTest(artifact=name): + self.assertGreater(os.path.getsize(os.path.join(d, name)), 0) + for s in live_s: + for kind in ("q", "k", "v", "control", "golden"): + name = f"{prefix}.S{s}.{kind}.bin" + with self.subTest(artifact=name): + self.assertGreater( + os.path.getsize(os.path.join(d, name)), 0 + ) + for kind in ("q", "k", "v", "control", "golden"): + name = f"sdpa_k16_qwen3.initial.{kind}.bin" + with self.subTest(artifact=name): + self.assertGreater(os.path.getsize(os.path.join(d, name)), 0) if __name__ == "__main__": diff --git a/backends/webgpu/test/ops/test_bitwise.py b/backends/webgpu/test/ops/test_bitwise.py index b61a37d9990..a8fee7a0095 100644 --- a/backends/webgpu/test/ops/test_bitwise.py +++ b/backends/webgpu/test/ops/test_bitwise.py @@ -19,6 +19,11 @@ import torch from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.backends.webgpu.test.ops.test_logical_and import ( + LOGICAL_BINARY_CASES, + logical_binary_gen_a, + logical_binary_gen_b, +) from executorch.exir import to_edge_transform_and_lower @@ -50,11 +55,10 @@ def g(shape): bw_gen_a = _bw_gen(0) -bw_gen_b = _bw_gen(1) # All shapes have numel % 4 == 0 (bool tensors pack 4 bytes/word). -SHAPES = [(4, 8), (2, 3, 8), (16, 16)] +BITWISE_NOT_SHAPES = ((4, 8), (2, 3, 8), (16, 16)) class BitwiseTest(unittest.TestCase): @@ -75,13 +79,17 @@ def _assert_delegates(self, mod, inputs, op_name, shape) -> None: ) def test_export_delegates(self) -> None: - for shape in SHAPES: - with self.subTest(shape=shape): - a = bw_gen_a(shape) - b = bw_gen_b(shape) + for case_name, shape in LOGICAL_BINARY_CASES: + with self.subTest(op="bitwise_and", case=case_name, shape=shape): + a = logical_binary_gen_a(shape) + b = logical_binary_gen_b(shape) self._assert_delegates( BitwiseAndModule(shape), (a, b), "bitwise_and", shape ) + + for shape in BITWISE_NOT_SHAPES: + with self.subTest(op="bitwise_not", shape=shape): + a = bw_gen_a(shape) self._assert_delegates( BitwiseNotModule(shape), (a,), "bitwise_not", shape ) diff --git a/backends/webgpu/test/ops/test_floor_divide.py b/backends/webgpu/test/ops/test_floor_divide.py index 6020480ff2c..141e8314cde 100644 --- a/backends/webgpu/test/ops/test_floor_divide.py +++ b/backends/webgpu/test/ops/test_floor_divide.py @@ -6,8 +6,8 @@ """`aten.div.Tensor_mode` module for the WebGPU op-test framework. -`FloorDivideModule` is imported by `cases.py`. Same-shape elementwise -`div(a, b, rounding_mode="floor")`. The kernel computes `floor(a/b)` mirroring +`FloorDivideModule` is imported by `cases.py`. Same-shape and broadcast +`div(a, b, rounding_mode="floor")` use `floor(a/b)`, mirroring the Vulkan `floor_divide` glsl (`floor(X/Y)`); this differs from torch's own fmod-corrected `div_floor` at rare fp boundaries, so the suite goldens against a `floor(a/b)` `golden_fn` (Vulkan-faithful), not this module's eager output. diff --git a/backends/webgpu/test/ops/test_logical_and.py b/backends/webgpu/test/ops/test_logical_and.py index 4c3939002d0..9cf474fd8de 100644 --- a/backends/webgpu/test/ops/test_logical_and.py +++ b/backends/webgpu/test/ops/test_logical_and.py @@ -4,17 +4,9 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -"""`aten.logical_and.default` module + configs for the WebGPU op-test framework. - -`LogicalAndModule` derives its two bool operands on-GPU from float inputs -(`a > 0`, `b > 0` via the delegated `gt.Tensor` against a baked zero buffer), so -the only runtime inputs are the two float tensors (the op-test framework is -float-input-only). `a`/`b` use distinct seeds so the two bool masks differ (each -~50% True, independent -> AND ~25% True), a real mix that a wrong op (e.g. OR) -would fail. Output is bool (byte-exact golden). `LogicalAndTest` is the -export-delegation smoke test. -""" +"""Delegation coverage for logical AND with packed truth-table inputs.""" +import math import unittest import torch @@ -32,29 +24,40 @@ def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: return torch.logical_and(a > self.z, b > self.z) -def _la_gen(seed): - # Distinct per-input seed so the two derived bool masks differ. - def g(shape): - gen = torch.Generator().manual_seed(seed) - return torch.randn(*shape, generator=gen, dtype=torch.float32) +LOGICAL_BINARY_CASES = ( + ("2d", (4, 8)), + ("3d", (2, 3, 8)), + ("sq", (16, 16)), + ("words63", (252,)), + ("words64", (256,)), + ("words65", (260,)), +) - return g +def _logical_binary_gen(pattern): + def generate(shape): + numel = math.prod(shape) + if numel == 0 or numel % len(pattern) != 0: + raise ValueError("logical-binary test shapes must have numel % 4 == 0") + return ( + torch.tensor(pattern, dtype=torch.float32) + .repeat(numel // len(pattern)) + .reshape(shape) + ) -la_gen_a = _la_gen(0) -la_gen_b = _la_gen(1) + return generate -# All shapes have numel % 4 == 0 (bool tensors pack 4 bytes/word). -SHAPES = [(4, 8), (2, 3, 8), (16, 16)] +logical_binary_gen_a = _logical_binary_gen((-1.0, -1.0, 1.0, 1.0)) +logical_binary_gen_b = _logical_binary_gen((-1.0, 1.0, -1.0, 1.0)) class LogicalAndTest(unittest.TestCase): def test_export_delegates(self) -> None: - for shape in SHAPES: - with self.subTest(shape=shape): - a = la_gen_a(shape) - b = la_gen_b(shape) + for case_name, shape in LOGICAL_BINARY_CASES: + with self.subTest(case=case_name, shape=shape): + a = logical_binary_gen_a(shape) + b = logical_binary_gen_b(shape) ep = torch.export.export(LogicalAndModule(shape).eval(), (a, b)) edge = to_edge_transform_and_lower( ep, partitioner=[VulkanPartitioner()] diff --git a/backends/webgpu/test/ops/test_logical_or.py b/backends/webgpu/test/ops/test_logical_or.py index d71c1493ea9..420fee26336 100644 --- a/backends/webgpu/test/ops/test_logical_or.py +++ b/backends/webgpu/test/ops/test_logical_or.py @@ -4,23 +4,18 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -"""`aten.logical_or.default` / `aten.bitwise_or.Tensor` (bool) modules + configs. - -Mirrors the logical_and/bitwise_and tests: the modules derive their two bool -operands on-GPU from float inputs (`a > 0`, `b > 0` via the delegated `gt.Tensor` -against a baked zero buffer), so the only runtime inputs are the two float -tensors (the op-test framework is float-input-only). `a`/`b` use distinct seeds -so the two bool masks differ (each ~50% True, independent -> OR ~75% True), a -real mix that a wrong op (e.g. AND) would fail. `bitwise_or` on bool is identical -to `logical_or` (shares the handler). Output is bool (byte-exact golden). -`LogicalOrTest` is the export-delegation smoke test. -""" +"""Delegation coverage for logical and bitwise OR truth-table inputs.""" import unittest import torch from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.backends.webgpu.test.ops.test_logical_and import ( + LOGICAL_BINARY_CASES, + logical_binary_gen_a, + logical_binary_gen_b, +) from executorch.exir import to_edge_transform_and_lower @@ -42,23 +37,6 @@ def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: return torch.bitwise_or(a > self.z, b > self.z) -def _lo_gen(seed): - # Distinct per-input seed so the two derived bool masks differ. - def g(shape): - gen = torch.Generator().manual_seed(seed) - return torch.randn(*shape, generator=gen, dtype=torch.float32) - - return g - - -lo_gen_a = _lo_gen(0) -lo_gen_b = _lo_gen(1) - - -# All shapes have numel % 4 == 0 (bool tensors pack 4 bytes/word). -SHAPES = [(4, 8), (2, 3, 8), (16, 16)] - - class LogicalOrTest(unittest.TestCase): def _assert_delegates(self, mod, inputs, op_name, shape) -> None: ep = torch.export.export(mod.eval(), inputs) @@ -77,10 +55,10 @@ def _assert_delegates(self, mod, inputs, op_name, shape) -> None: ) def test_export_delegates(self) -> None: - for shape in SHAPES: - with self.subTest(shape=shape): - a = lo_gen_a(shape) - b = lo_gen_b(shape) + for case_name, shape in LOGICAL_BINARY_CASES: + with self.subTest(case=case_name, shape=shape): + a = logical_binary_gen_a(shape) + b = logical_binary_gen_b(shape) self._assert_delegates( LogicalOrModule(shape), (a, b), "logical_or", shape ) diff --git a/backends/webgpu/test/ops/test_minimum.py b/backends/webgpu/test/ops/test_minimum.py index f8f8088fef5..ba466317d5b 100644 --- a/backends/webgpu/test/ops/test_minimum.py +++ b/backends/webgpu/test/ops/test_minimum.py @@ -6,8 +6,8 @@ """`aten.minimum.default` module for the WebGPU op-test framework. -`MinimumModule` is imported by `cases.py`. minimum is a same-shape elementwise -binary op mirroring the landed `add`/`mul` pattern (flat 2D-dispatch kernel). +`MinimumModule` is imported by `cases.py`. minimum is an elementwise binary op +with the same-shape fast path and broadcast indexing used by `mul`. """ import torch diff --git a/backends/webgpu/test/ops/test_pow.py b/backends/webgpu/test/ops/test_pow.py index 93596dc5d9f..8498bd4f7aa 100644 --- a/backends/webgpu/test/ops/test_pow.py +++ b/backends/webgpu/test/ops/test_pow.py @@ -6,8 +6,8 @@ """`aten.pow.Tensor_Tensor` module for the WebGPU op-test framework. -`PowModule` is imported by `cases.py`. Same-shape elementwise `pow(a, b)`; the -suite uses a POSITIVE base so `pow(neg, frac)` (NaN) is never exercised. +`PowModule` is imported by `cases.py`. The suite covers same-shape and broadcast +`pow(a, b)` with POSITIVE bases so `pow(neg, frac)` never produces NaN. """ import torch diff --git a/backends/webgpu/test/ops/test_quantized_linear.py b/backends/webgpu/test/ops/test_quantized_linear.py index 54040a54bd7..c68553c57ca 100644 --- a/backends/webgpu/test/ops/test_quantized_linear.py +++ b/backends/webgpu/test/ops/test_quantized_linear.py @@ -15,6 +15,9 @@ CONFIGS table and reconstructs the identical deterministic ramp input bit-for-bit. """ +import hashlib +import hmac +import json import os import unittest from dataclasses import dataclass @@ -22,8 +25,10 @@ import numpy as np import torch -from executorch.backends.vulkan import VulkanPartitioner -from executorch.exir import to_edge_transform_and_lower +from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.exir import ExecutorchBackendConfig, to_edge_transform_and_lower +from executorch.exir.backend.utils import get_delegates, get_non_lowered_nodes +from executorch.exir.passes import MemoryPlanningPass from torchao.quantization.granularity import PerGroup from torchao.quantization.quant_api import IntxWeightOnlyConfig, quantize_ @@ -89,8 +94,10 @@ class Q4gswConfig: ] -def _make_quantized_model(k: int, n: int, group_size: int) -> torch.nn.Module: - torch.manual_seed(0) # load-bearing: fixes the weights the golden derives from +def _make_quantized_model( + k: int, n: int, group_size: int, seed: int = 0 +) -> torch.nn.Module: + torch.manual_seed(seed) # load-bearing: fixes the weights used by the golden m = torch.nn.Linear(k, n, bias=False).eval() quantize_( m, @@ -127,6 +134,149 @@ def _export(m: torch.nn.Module, x: torch.Tensor): ).to_executorch() +OUTPUT_SUPPRESSION_K = 64 +OUTPUT_SUPPRESSION_N = 64 +OUTPUT_SUPPRESSION_GROUP = 32 + + +class _DirectFinalQ4(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.q4 = _make_quantized_model( + OUTPUT_SUPPRESSION_K, + OUTPUT_SUPPRESSION_N, + OUTPUT_SUPPRESSION_GROUP, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.q4(x) + + +class _Q4ThenAdd(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.q4 = _make_quantized_model( + OUTPUT_SUPPRESSION_K, + OUTPUT_SUPPRESSION_N, + OUTPUT_SUPPRESSION_GROUP, + ) + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + q4 = self.q4(x) + return q4, q4 + q4 + + +class _UnrelatedThenFinalQ4(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.q4 = _make_quantized_model( + OUTPUT_SUPPRESSION_K, + OUTPUT_SUPPRESSION_N, + OUTPUT_SUPPRESSION_GROUP, + ) + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + unrelated = x + x + return unrelated, self.q4(unrelated) + + +class _DuplicateFinalQ4(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.q4 = _make_quantized_model( + OUTPUT_SUPPRESSION_K, + OUTPUT_SUPPRESSION_N, + OUTPUT_SUPPRESSION_GROUP, + ) + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + q4 = self.q4(x) + return q4, q4 + + +class _Q4ThenPortableSum(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.q4 = _make_quantized_model( + OUTPUT_SUPPRESSION_K, + OUTPUT_SUPPRESSION_N, + OUTPUT_SUPPRESSION_GROUP, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.q4(x).sum() + + +def _output_suppression_cases() -> list[tuple[str, torch.nn.Module]]: + return [ + ("direct_final_q4", _DirectFinalQ4()), + ("q4_then_add", _Q4ThenAdd()), + ("unrelated_then_final_q4", _UnrelatedThenFinalQ4()), + ] + + +def _verify_output_suppression_certificate( + certificate: dict[str, object], pte: bytes, method: str = "forward" +) -> None: + if certificate.get("schema") != 1: + raise ValueError("invalid output-suppression certificate schema") + if certificate.get("method") != method: + raise ValueError("output-suppression certificate method mismatch") + if certificate.get("delegate_count") != 1: + raise ValueError("output suppression requires exactly one delegate") + if certificate.get("portable_node_count") != 0: + raise ValueError("output suppression rejects portable nodes") + outputs = certificate.get("method_output_nodes") + if not isinstance(outputs, list) or not outputs: + raise ValueError("output-suppression certificate has no method outputs") + if len(set(outputs)) != len(outputs): + raise ValueError("aliased method outputs cannot be suppressed") + actual_hash = hashlib.sha256(pte).hexdigest() + expected_hash = certificate.get("pte_sha256") + if not isinstance(expected_hash, str) or not hmac.compare_digest( + actual_hash, expected_hash + ): + raise ValueError("output-suppression certificate PTE hash mismatch") + + +def _lower_and_certify(model: torch.nn.Module, x: torch.Tensor): + edge = to_edge_transform_and_lower( + torch.export.export(model, (x,)), partitioner=[VulkanPartitioner()] + ) + graph = edge.exported_program().graph_module.graph + delegates = get_delegates(graph) + non_lowered = get_non_lowered_nodes(graph) + output_node = next(node for node in graph.nodes if node.op == "output") + method_outputs = output_node.args[0] + if not isinstance(method_outputs, (tuple, list)): + method_outputs = (method_outputs,) + if len(delegates) != 1: + raise ValueError(f"expected one delegate, got {len(delegates)}") + if non_lowered: + raise ValueError(f"non-lowered nodes: {non_lowered}") + if len({id(output) for output in method_outputs}) != len(method_outputs): + raise ValueError("aliased method outputs cannot be suppressed") + et = edge.to_executorch( + ExecutorchBackendConfig( + memory_planning_pass=MemoryPlanningPass(alloc_graph_output=False) + ) + ) + certificate = { + "schema": 1, + "method": "forward", + "pte_sha256": hashlib.sha256(et.buffer).hexdigest(), + "delegate_count": len(delegates), + "portable_node_count": len(non_lowered), + "method_output_nodes": [output.name for output in method_outputs], + } + _verify_output_suppression_certificate(certificate, et.buffer) + return et, certificate + + +def _lower_fully_delegated(model: torch.nn.Module, x: torch.Tensor): + return _lower_and_certify(model, x)[0] + + class TestQuantizedLinear(unittest.TestCase): def test_export_delegates(self) -> None: # Each (non-heavy) config must fuse to a VulkanBackend delegate (q4gsw); @@ -144,6 +294,37 @@ def test_export_delegates(self) -> None: ) self.assertTrue(found, f"no VulkanBackend delegate in {cfg.name}") + def test_output_suppression_models_are_fully_delegated(self) -> None: + x = _ramp_input(1, OUTPUT_SUPPRESSION_K) + for name, model in _output_suppression_cases(): + with self.subTest(case=name): + et = _lower_fully_delegated(model, x) + delegate_ids = [ + delegate.id + for plan in et.executorch_program.execution_plan + for delegate in plan.delegates + ] + self.assertEqual(delegate_ids, ["VulkanBackend"]) + + def test_output_suppression_rejects_aliased_method_outputs(self) -> None: + x = _ramp_input(1, OUTPUT_SUPPRESSION_K) + with self.assertRaisesRegex(ValueError, "aliased method outputs"): + _lower_fully_delegated(_DuplicateFinalQ4(), x) + + def test_output_suppression_rejects_portable_continuation(self) -> None: + x = _ramp_input(1, OUTPUT_SUPPRESSION_K) + with self.assertRaisesRegex(ValueError, "non-lowered nodes"): + _lower_fully_delegated(_Q4ThenPortableSum(), x) + + def test_output_suppression_certificate_binds_exact_pte(self) -> None: + x = _ramp_input(1, OUTPUT_SUPPRESSION_K) + et, certificate = _lower_and_certify(_DirectFinalQ4(), x) + _verify_output_suppression_certificate(certificate, et.buffer) + with self.assertRaises(ValueError): + _verify_output_suppression_certificate( + certificate, et.buffer + b"different" + ) + def test_golden_matches_eager(self) -> None: # Dual oracle (mirrors SDPA test_golden_matches_eager_op): the fp64 dequant- # matmul truth and torchao's own fp32 quantized forward are independent refs @@ -189,5 +370,37 @@ def export_all_quantized_linear_models( export_quantized_linear_model(cfg, pte, golden) +def export_output_suppression_models(out_dir: str) -> None: + os.makedirs(out_dir, exist_ok=True) + x = _ramp_input(1, OUTPUT_SUPPRESSION_K) + x.numpy().astype(" torch.Tensor: return ((flat % 17) - 8).div(16.0).reshape(shape) +# Shared structural and manifest-driven extrema case authority. +EXTREMA_CONFIGS = ( + ("keepdim_2d", (37, 41), -1, True, "default"), + ("nodim_2d", (37, 41), -1, False, "default"), + ("keepdim_3d", (5, 7, 11), -1, True, "default"), + ("nodim_3d", (5, 7, 11), -1, False, "default"), + ("sign_trap_63_drop", (3, 63), -1, False, "sign_trap"), + ("tie_64_keep", (2, 64), -1, True, "tie"), + ("tie_65_posdim_drop", (2, 65), 1, False, "tie"), + ("sign_trap_255_keep", (2, 255), -1, True, "sign_trap"), + ("tie_256_drop", (2, 256), -1, False, "tie"), + ("tie_257_posdim_keep", (2, 257), 1, True, "tie"), +) + + +def _sign_trap_input(shape, *, for_max: bool) -> torch.Tensor: + magnitude = ((torch.arange(math.prod(shape), dtype=torch.float32) % 31) + 1).div( + 8.0 + ) + rows = magnitude.reshape(-1, shape[-1]) + for row_index, row in enumerate(rows): + row.add_(float(row_index)) + values = -magnitude if for_max else magnitude + return values.reshape(shape) + + +def amax_sign_trap_input(shape) -> torch.Tensor: + return _sign_trap_input(shape, for_max=True) + + +def amin_sign_trap_input(shape) -> torch.Tensor: + return _sign_trap_input(shape, for_max=False) + + +def _tie_input(shape, *, for_max: bool) -> torch.Tensor: + values = ((torch.arange(math.prod(shape), dtype=torch.float32) % 29) - 14).div(8.0) + rows = values.reshape(-1, shape[-1]) + for row_index, row in enumerate(rows): + extreme = 16.0 + float(row_index) + if not for_max: + extreme = -extreme + row[-2] = extreme + row[-1] = extreme + return values.reshape(shape) + + +def amax_tie_input(shape) -> torch.Tensor: + return _tie_input(shape, for_max=True) + + +def amin_tie_input(shape) -> torch.Tensor: + return _tie_input(shape, for_max=False) + + +def _extrema_input(op: str, input_class: str, shape) -> torch.Tensor: + if input_class == "default": + return _det_input(shape) + if input_class == "sign_trap": + return ( + amax_sign_trap_input(shape) if op == "amax" else amin_sign_trap_input(shape) + ) + return amax_tie_input(shape) if op == "amax" else amin_tie_input(shape) + + def _export(m: torch.nn.Module, x: torch.Tensor): ep = torch.export.export(m, (x,)) return to_edge_transform_and_lower( @@ -128,21 +194,55 @@ def export_reduce_model( class AmaxModule(torch.nn.Module): - def __init__(self, keepdim: bool) -> None: + def __init__(self, keepdim: bool, dim: int = -1) -> None: super().__init__() self.keepdim = keepdim + self.dim = dim def forward(self, x: torch.Tensor) -> torch.Tensor: - return torch.amax(x, dim=-1, keepdim=self.keepdim) + return torch.amax(x, dim=self.dim, keepdim=self.keepdim) class AminModule(torch.nn.Module): - def __init__(self, keepdim: bool) -> None: + def __init__(self, keepdim: bool, dim: int = -1) -> None: super().__init__() self.keepdim = keepdim + self.dim = dim def forward(self, x: torch.Tensor) -> torch.Tensor: - return torch.amin(x, dim=-1, keepdim=self.keepdim) + return torch.amin(x, dim=self.dim, keepdim=self.keepdim) + + +class TestExtrema(unittest.TestCase): + def test_config_contract(self) -> None: + self.assertEqual( + EXTREMA_CONFIGS, + ( + ("keepdim_2d", (37, 41), -1, True, "default"), + ("nodim_2d", (37, 41), -1, False, "default"), + ("keepdim_3d", (5, 7, 11), -1, True, "default"), + ("nodim_3d", (5, 7, 11), -1, False, "default"), + ("sign_trap_63_drop", (3, 63), -1, False, "sign_trap"), + ("tie_64_keep", (2, 64), -1, True, "tie"), + ("tie_65_posdim_drop", (2, 65), 1, False, "tie"), + ("sign_trap_255_keep", (2, 255), -1, True, "sign_trap"), + ("tie_256_drop", (2, 256), -1, False, "tie"), + ("tie_257_posdim_keep", (2, 257), 1, True, "tie"), + ), + ) + + def test_exports_fully_delegated(self) -> None: + for op, module_cls in (("amax", AmaxModule), ("amin", AminModule)): + for name, shape, dim, keepdim, input_class in EXTREMA_CONFIGS: + with self.subTest(op=op, config=name): + x = _extrema_input(op, input_class, shape) + ep = torch.export.export(module_cls(keepdim, dim).eval(), (x,)) + edge = to_edge_transform_and_lower( + ep, partitioner=[VulkanPartitioner()] + ) + graph = edge.exported_program().graph_module.graph + self.assertEqual(len(get_delegates(graph)), 1) + self.assertEqual(get_non_lowered_nodes(graph), []) if __name__ == "__main__": diff --git a/backends/webgpu/test/ops/test_rope_hf.py b/backends/webgpu/test/ops/test_rope_hf.py index 6e26fe53b17..f1929ac8d3b 100644 --- a/backends/webgpu/test/ops/test_rope_hf.py +++ b/backends/webgpu/test/ops/test_rope_hf.py @@ -19,6 +19,7 @@ compare (it has no ATen). """ +import os import unittest from collections import namedtuple @@ -26,8 +27,12 @@ import torch from executorch.backends.vulkan import VulkanPartitioner -from executorch.examples.models.llama.rope import hf_apply_rotary_emb +from executorch.examples.models.llama.rope import ( + hf_apply_rotary_emb, + hf_precompute_freqs_cis, +) from executorch.exir import to_edge_transform_and_lower +from executorch.exir.backend.utils import get_delegates, get_non_lowered_nodes # B batch, S tokens, NH query heads, NKV kv heads (NH != NKV so the two outputs # are distinguishable by numel), HD head dim (even; full rotary, rotary_dim==HD). @@ -39,6 +44,20 @@ Shape("decode", 1, 1, 16, 8, 128), ] +DYNAMIC_BATCH = 1 +DYNAMIC_SEQ = 1 +DYNAMIC_N_HEADS_Q = 16 +DYNAMIC_N_HEADS_K = 8 +DYNAMIC_HEAD_DIM = 128 +DYNAMIC_MAX_SEQ = 16 +DYNAMIC_POSITIONS = (0, 7, 15) +DYNAMIC_SEQUENCE_CASES = ( + (DYNAMIC_MAX_SEQ, 0), + (5, 7), + (1, DYNAMIC_MAX_SEQ - 1), + (DYNAMIC_MAX_SEQ, 0), +) + class HfRope(torch.nn.Module): # unsqueeze_dim=1: freqs [S, HD] -> [S, 1, HD] broadcasts over (B, NH) of the @@ -47,6 +66,26 @@ def forward(self, xq, xk, freqs_cos, freqs_sin): return hf_apply_rotary_emb(xq, xk, freqs_cos, freqs_sin, unsqueeze_dim=1) +class DynamicHfRope(torch.nn.Module): + def forward( + self, + xq: torch.Tensor, + xk: torch.Tensor, + freqs_cos: torch.Tensor, + freqs_sin: torch.Tensor, + input_pos: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + start_pos = input_pos[0].item() + torch._check_is_size(start_pos) + torch._check(start_pos + xq.shape[1] <= freqs_cos.shape[0]) + return hf_apply_rotary_emb( + xq, + xk, + freqs_cos.narrow(0, start_pos, xq.shape[1]), + freqs_sin.narrow(0, start_pos, xq.shape[1]), + ) + + def _ramp(numel: int, mod: int, off: int) -> torch.Tensor: # ((i % mod) - off) / 16: exact in fp32, matches test_webgpu_native.cpp. idx = torch.arange(numel, dtype=torch.int64) @@ -68,6 +107,32 @@ def _inputs( return xq, xk, freqs_cos, freqs_sin +def _dynamic_inputs(seq: int = DYNAMIC_SEQ) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, +]: + xq = _ramp( + DYNAMIC_BATCH * seq * DYNAMIC_N_HEADS_Q * DYNAMIC_HEAD_DIM, + 17, + 8, + ).reshape(DYNAMIC_BATCH, seq, DYNAMIC_N_HEADS_Q, DYNAMIC_HEAD_DIM) + xk = _ramp( + DYNAMIC_BATCH * seq * DYNAMIC_N_HEADS_K * DYNAMIC_HEAD_DIM, + 13, + 6, + ).reshape(DYNAMIC_BATCH, seq, DYNAMIC_N_HEADS_K, DYNAMIC_HEAD_DIM) + freqs_cos, freqs_sin = hf_precompute_freqs_cis( + DYNAMIC_HEAD_DIM, + DYNAMIC_MAX_SEQ, + theta=10000.0, + ) + input_pos = torch.tensor([0], dtype=torch.long) + return xq, xk, freqs_cos, freqs_sin, input_pos + + def _golden( xq: torch.Tensor, xk: torch.Tensor, @@ -78,27 +143,123 @@ def _golden( return torch.ops.et_vk.apply_rotary_emb_hf.default(xq, xk, freqs_cos, freqs_sin, 0) -def _export(inputs): +def _dynamic_golden( + xq: torch.Tensor, + xk: torch.Tensor, + freqs_cos: torch.Tensor, + freqs_sin: torch.Tensor, + position: int, +) -> tuple[torch.Tensor, torch.Tensor]: + return hf_apply_rotary_emb( + xq, + xk, + freqs_cos[position : position + xq.shape[1]], + freqs_sin[position : position + xq.shape[1]], + ) + + +def _assert_fully_delegated(edge) -> None: + graph = edge.exported_program().graph_module.graph + delegates = get_delegates(graph) + portable = get_non_lowered_nodes(graph) + if len(delegates) != 1: + raise AssertionError(f"expected one delegate, got {len(delegates)}") + if portable: + raise AssertionError(f"unexpected non-lowered nodes: {portable}") + + +def _lower(inputs): ep = torch.export.export(HfRope().eval(), inputs) - return to_edge_transform_and_lower( - ep, partitioner=[VulkanPartitioner()] - ).to_executorch() + edge = to_edge_transform_and_lower(ep, partitioner=[VulkanPartitioner()]) + _assert_fully_delegated(edge) + return edge + + +def _export(inputs): + return _lower(inputs).to_executorch() + + +def _lower_dynamic_program(): + inputs = _dynamic_inputs() + with torch._dynamo.config.patch(capture_scalar_outputs=True): + ep = torch.export.export(DynamicHfRope().eval(), inputs) + + symints = [ + node + for node in ep.graph_module.graph.nodes + if isinstance(node.meta.get("val"), torch.SymInt) + ] + if not symints: + raise AssertionError("input_pos did not lower to a SymInt") + + edge = to_edge_transform_and_lower(ep, partitioner=[VulkanPartitioner()]) + _assert_fully_delegated(edge) + return edge + + +def _lower_dynamic_sequence_program(): + inputs = _dynamic_inputs(DYNAMIC_MAX_SEQ) + s_dim = torch.export.Dim("rope_hf_s", min=1, max=DYNAMIC_MAX_SEQ) + dynamic_shapes = ({1: s_dim}, {1: s_dim}, None, None, None) + with torch._dynamo.config.patch(capture_scalar_outputs=True): + ep = torch.export.export( + DynamicHfRope().eval(), + inputs, + dynamic_shapes=dynamic_shapes, + ) + + scalar_symints = [ + node + for node in ep.graph_module.graph.nodes + if isinstance(node.meta.get("val"), torch.SymInt) + ] + if not scalar_symints: + raise AssertionError("input_pos did not lower to a SymInt") + xq_placeholder = next( + node + for node in ep.graph_module.graph.nodes + if node.op == "placeholder" and node.target == "xq" + ) + if not isinstance(xq_placeholder.meta["val"].shape[1], torch.SymInt): + raise AssertionError("query sequence dimension did not remain symbolic") + + edge = to_edge_transform_and_lower(ep, partitioner=[VulkanPartitioner()]) + _assert_fully_delegated(edge) + return edge + + +def _export_dynamic_program(): + edge = _lower_dynamic_program() + + et = edge.to_executorch() + delegate_ids = [ + delegate.id + for plan in et.executorch_program.execution_plan + for delegate in plan.delegates + ] + if delegate_ids != ["VulkanBackend"]: + raise AssertionError(f"unexpected delegates: {delegate_ids}") + return et + + +def _export_dynamic_sequence_program(): + edge = _lower_dynamic_sequence_program() + et = edge.to_executorch() + delegate_ids = [ + delegate.id + for plan in et.executorch_program.execution_plan + for delegate in plan.delegates + ] + if delegate_ids != ["VulkanBackend"]: + raise AssertionError(f"unexpected delegates: {delegate_ids}") + return et class TestRopeHf(unittest.TestCase): def test_export_delegates(self) -> None: for shape in SHAPES: with self.subTest(shape=shape.name): - et = _export(_inputs(shape)) - found = any( - d.id == "VulkanBackend" - for plan in et.executorch_program.execution_plan - for d in plan.delegates - ) - self.assertTrue( - found, - "Expected a VulkanBackend delegate (apply_rotary_emb_hf " "fusion)", - ) + self.assertIsNotNone(_lower(_inputs(shape))) def test_golden_matches_eager(self) -> None: # The et_vk golden must equal the real HF rotate-half apply_rotary_emb, @@ -112,6 +273,44 @@ def test_golden_matches_eager(self) -> None: torch.testing.assert_close(gq, eq, atol=1e-5, rtol=1e-5) torch.testing.assert_close(gk, ek, atol=1e-5, rtol=1e-5) + def test_dynamic_export_is_fully_delegated(self) -> None: + self.assertIsNotNone(_lower_dynamic_program()) + + def test_dynamic_position_goldens_match_custom_op(self) -> None: + xq, xk, freqs_cos, freqs_sin, _ = _dynamic_inputs() + self.assertNotEqual(xq.shape[2], xk.shape[2]) + position_outputs = [] + for position in DYNAMIC_POSITIONS: + with self.subTest(position=position): + expected_q, expected_k = _dynamic_golden( + xq, xk, freqs_cos, freqs_sin, position + ) + position_outputs.append(expected_q) + actual_q, actual_k = torch.ops.et_vk.apply_rotary_emb_hf.default( + xq, xk, freqs_cos, freqs_sin, position + ) + torch.testing.assert_close(actual_q, expected_q) + torch.testing.assert_close(actual_k, expected_k) + self.assertFalse(torch.allclose(position_outputs[0], position_outputs[1])) + self.assertFalse(torch.allclose(position_outputs[1], position_outputs[2])) + + def test_dynamic_sequence_export_is_fully_delegated(self) -> None: + self.assertIsNotNone(_lower_dynamic_sequence_program()) + + def test_dynamic_sequence_goldens_match_custom_op(self) -> None: + _, _, freqs_cos, freqs_sin, _ = _dynamic_inputs(DYNAMIC_MAX_SEQ) + for seq, position in dict.fromkeys(DYNAMIC_SEQUENCE_CASES): + with self.subTest(seq=seq, position=position): + xq, xk, _, _, _ = _dynamic_inputs(seq) + expected_q, expected_k = _dynamic_golden( + xq, xk, freqs_cos, freqs_sin, position + ) + actual_q, actual_k = torch.ops.et_vk.apply_rotary_emb_hf.default( + xq, xk, freqs_cos, freqs_sin, position + ) + torch.testing.assert_close(actual_q, expected_q) + torch.testing.assert_close(actual_k, expected_k) + def export_rope_hf_model( pte_path: str, xq_golden_path: str, xk_golden_path: str, shape_name: str = "multi" @@ -132,5 +331,58 @@ def export_rope_hf_model( ) +def export_rope_hf_dynamic(out_dir: str) -> None: + os.makedirs(out_dir, exist_ok=True) + xq, xk, freqs_cos, freqs_sin, _ = _dynamic_inputs() + et = _export_dynamic_program() + with open(os.path.join(out_dir, "rope_hf_dynamic.pte"), "wb") as output: + output.write(et.buffer) + + for name, tensor in ( + ("xq", xq), + ("xk", xk), + ("freqs_cos", freqs_cos), + ("freqs_sin", freqs_sin), + ): + tensor.detach().numpy().astype(" None: + os.makedirs(out_dir, exist_ok=True) + _, _, freqs_cos, freqs_sin, _ = _dynamic_inputs(DYNAMIC_MAX_SEQ) + et = _export_dynamic_sequence_program() + prefix = "rope_hf_dynamic_sequence" + with open(os.path.join(out_dir, f"{prefix}.pte"), "wb") as output: + output.write(et.buffer) + + for name, tensor in (("freqs_cos", freqs_cos), ("freqs_sin", freqs_sin)): + tensor.detach().numpy().astype(" large logits (softmax stress) + kv_f16: bool = False # Single source of truth, mirrored by the C++ CONFIGS table in the native test. @@ -64,6 +65,10 @@ class SdpaConfig: # 2D-dispatch cap (>65535 wg): S=512 folds QK; S=2048 folds QK+softmax+AV (cap+1). SdpaConfig("llama1b_prefill_512", 32, 8, 64, 512, 512, 0), SdpaConfig("llama1b_prefill_2048", 32, 8, 64, 2048, 2048, 0), + # denom=10 intentionally makes K/V values lossy in fp16, so native + # execution exercises the real fp32->fp16->fp32 cache conversion path. + SdpaConfig("qwen3_prefill", 16, 8, 128, 128, 256, 0, 10.0, kv_f16=True), + SdpaConfig("qwen3_odd_boundary", 16, 8, 128, 17, 64, 31, 10.0, kv_f16=True), ] @@ -83,6 +88,7 @@ class ReplaySeq: d: int # head dim cmax: int # kv-cache capacity (>= sum(seq_lens)) seq_lens: tuple[int, ...] + kv_f16: bool = False # Mirror Vulkan sdpa_test.cpp:856/867/875 (3 param sets); cmax = sum rounded up. @@ -90,11 +96,18 @@ class ReplaySeq: ReplaySeq("small", 8, 4, 4, 16, (3, 1, 1, 5, 1, 1, 2)), ReplaySeq("small_d", 6, 2, 8, 16, (3, 1, 1, 5, 1, 1)), ReplaySeq("llama3", 24, 8, 128, 256, (111, 1, 1, 1, 57, 1, 1)), + ReplaySeq("qwen3_fd", 16, 8, 128, 64, (17, 1), kv_f16=True), ] +DYNAMIC_REPLAY_SEQS = [seq for seq in REPLAY_SEQS if not seq.kv_f16] -# (head_dim, num_heads, num_kv_heads) from sdpa_test.cpp:856/867/875 -- guards a -# transposition of the (hq, hkv, d) field order against the Vulkan source. -VULKAN_PARAMS = {"small": (4, 8, 4), "small_d": (8, 6, 2), "llama3": (128, 24, 8)} +# Guards transposition of the (hq, hkv, d) field order. The first three values +# mirror Vulkan sdpa_test.cpp:856/867/875; Qwen3 extends the same contract. +VULKAN_PARAMS = { + "small": (4, 8, 4), + "small_d": (8, 6, 2), + "llama3": (128, 24, 8), + "qwen3_fd": (128, 16, 8), +} class SdpaModule(torch.nn.Module): @@ -178,6 +191,12 @@ def _det_inputs(cfg: SdpaConfig): return q, k, v, k_cache, v_cache +def _round_kv_for_storage(cfg: SdpaConfig, *tensors: torch.Tensor): + if not cfg.kv_f16: + return tensors + return tuple(tensor.to(torch.float16).to(torch.float32) for tensor in tensors) + + def _golden(cfg: SdpaConfig, q, k, v, k_cache, v_cache) -> torch.Tensor: """Reference attention output [1,S,Hq,D], computed in fp64 then cast to fp32. @@ -189,6 +208,7 @@ def _golden(cfg: SdpaConfig, q, k, v, k_cache, v_cache) -> torch.Tensor: """ context_len = cfg.s + cfg.input_pos g = cfg.hq // cfg.hkv + k, v, k_cache, v_cache = _round_kv_for_storage(cfg, k, v, k_cache, v_cache) qd, kd, vd = q.double(), k.double(), v.double() kcd, vcd = k_cache.double(), v_cache.double() @@ -228,6 +248,50 @@ def _export_pte(cfg: SdpaConfig, q, k, v, kc, vc): class TestSdpa(unittest.TestCase): + def test_qwen3_fixture_contract(self) -> None: + configs = {cfg.name: cfg for cfg in CONFIGS} + expected_geometries = { + "qwen3_prefill": (16, 8, 128, 128, 256, 0), + "qwen3_odd_boundary": (16, 8, 128, 17, 64, 31), + } + for name, geometry in expected_geometries.items(): + with self.subTest(config=name): + self.assertIn(name, configs) + cfg = configs[name] + self.assertEqual( + (cfg.hq, cfg.hkv, cfg.d, cfg.s, cfg.cmax, cfg.input_pos), + geometry, + ) + self.assertTrue(cfg.kv_f16) + + replays = {seq.name: seq for seq in REPLAY_SEQS} + self.assertIn("qwen3_fd", replays) + qwen3_fd = replays["qwen3_fd"] + self.assertEqual((qwen3_fd.hq, qwen3_fd.hkv, qwen3_fd.d), (16, 8, 128)) + self.assertEqual(qwen3_fd.seq_lens, (17, 1)) + self.assertTrue(qwen3_fd.kv_f16) + + probe = torch.tensor([0.1], dtype=torch.float32) + (rounded,) = _round_kv_for_storage(configs["qwen3_prefill"], probe) + expected = probe.to(torch.float16).to(torch.float32) + torch.testing.assert_close(rounded, expected, atol=0.0, rtol=0.0) + self.assertFalse(torch.equal(rounded, probe)) + + boundary = configs["qwen3_odd_boundary"] + q, k, v, k_cache, v_cache = _det_inputs(boundary) + self.assertGreater(torch.count_nonzero(k_cache).item(), 0) + self.assertGreater(torch.count_nonzero(v_cache).item(), 0) + initialized = _golden(boundary, q, k, v, k_cache, v_cache) + cleared = _golden( + boundary, + q, + k, + v, + torch.zeros_like(k_cache), + torch.zeros_like(v_cache), + ) + self.assertFalse(torch.equal(initialized, cleared)) + def test_sdpa_export_delegates(self) -> None: for cfg in CONFIGS: with self.subTest(config=cfg.name): @@ -248,7 +312,12 @@ def test_golden_matches_eager_op(self) -> None: for cfg in CONFIGS: with self.subTest(config=cfg.name): q, k, v, kc, vc = _det_inputs(cfg) - eager = SdpaModule(cfg.input_pos)(q, k, v, kc.clone(), vc.clone()) + eager_k, eager_v, eager_kc, eager_vc = _round_kv_for_storage( + cfg, k, v, kc, vc + ) + eager = SdpaModule(cfg.input_pos)( + q, eager_k, eager_v, eager_kc.clone(), eager_vc.clone() + ) golden = _golden(cfg, q, k, v, kc, vc) torch.testing.assert_close(eager, golden, atol=1e-4, rtol=1e-4) @@ -277,10 +346,16 @@ def test_replay_golden_matches_eager(self) -> None: s, seq.cmax, input_pos, + kv_f16=seq.kv_f16, ) q, k, v = _step_inputs(seq, t, s) golden = _golden(cfg, q, k, v, kc, vc) - eager = SdpaModule(input_pos)(q, k, v, kc.clone(), vc.clone()) + eager_k, eager_v, eager_kc, eager_vc = _round_kv_for_storage( + cfg, k, v, kc, vc + ) + eager = SdpaModule(input_pos)( + q, eager_k, eager_v, eager_kc.clone(), eager_vc.clone() + ) torch.testing.assert_close(eager, golden, atol=1e-4, rtol=1e-4) kc[0, input_pos : input_pos + s] = k[0] vc[0, input_pos : input_pos + s] = v[0] @@ -302,6 +377,7 @@ def test_replay_export_delegates(self) -> None: s, seq.cmax, input_pos, + kv_f16=seq.kv_f16, ) q, k, v = _step_inputs(seq, t, s) et = _export_pte(cfg, q, k, v, kc, vc) @@ -347,7 +423,14 @@ def export_replay_sequences(out_dir: str) -> None: input_pos = 0 for t, s in enumerate(seq.seq_lens): cfg = SdpaConfig( - f"{seq.name}_step{t}", seq.hq, seq.hkv, seq.d, s, seq.cmax, input_pos + f"{seq.name}_step{t}", + seq.hq, + seq.hkv, + seq.d, + s, + seq.cmax, + input_pos, + kv_f16=seq.kv_f16, ) q, k, v = _step_inputs(seq, t, s) et = _export_pte(cfg, q, k, v, ref_kc, ref_vc) @@ -421,7 +504,7 @@ def export_dynamic_decode(out_dir: str) -> None: Mirrors the host accumulation the native test threads: at step t the golden attends over input_pos=t prior tokens plus the new token. """ - for seq in REPLAY_SEQS: + for seq in DYNAMIC_REPLAY_SEQS: assert DYN_DECODE_STEPS <= seq.cmax, f"{seq.name}: decode exceeds cmax" et = _export_dyn_pte(seq, 1) pte_path = os.path.join(out_dir, f"sdpa_dyn_{seq.name}.pte") @@ -431,7 +514,14 @@ def export_dynamic_decode(out_dir: str) -> None: ref_vc = torch.zeros(1, seq.cmax, seq.hkv, seq.d) for t in range(DYN_DECODE_STEPS): cfg = SdpaConfig( - f"dyn_{seq.name}_step{t}", seq.hq, seq.hkv, seq.d, 1, seq.cmax, t + f"dyn_{seq.name}_step{t}", + seq.hq, + seq.hkv, + seq.d, + 1, + seq.cmax, + t, + kv_f16=seq.kv_f16, ) q, k, v = _step_inputs(seq, t, 1) golden = _golden(cfg, q, k, v, ref_kc, ref_vc).numpy().astype(" None: def test_dynamic_decode_golden_matches_eager(self) -> None: # The threaded-cache decode golden must equal the eager op step-by-step. - for seq in REPLAY_SEQS: + for seq in DYNAMIC_REPLAY_SEQS: ref_kc = torch.zeros(1, seq.cmax, seq.hkv, seq.d) ref_vc = torch.zeros(1, seq.cmax, seq.hkv, seq.d) for t in range(DYN_DECODE_STEPS): cfg = SdpaConfig( - f"dyn_{seq.name}_step{t}", seq.hq, seq.hkv, seq.d, 1, seq.cmax, t + f"dyn_{seq.name}_step{t}", + seq.hq, + seq.hkv, + seq.d, + 1, + seq.cmax, + t, + kv_f16=seq.kv_f16, ) q, k, v = _step_inputs(seq, t, 1) golden = _golden(cfg, q, k, v, ref_kc, ref_vc) @@ -499,7 +596,7 @@ def export_incache_decode(out_dir: str) -> None: """One sdpa_incache_.pte (mutable-buffer KV cache) + per-step decode goldens. forward() feeds only q/k/v + input_pos; the cache persists in-graph. """ - for seq in REPLAY_SEQS: + for seq in DYNAMIC_REPLAY_SEQS: assert DYN_DECODE_STEPS <= seq.cmax, f"{seq.name}: decode exceeds cmax" m = DecodeCacheModule(seq.hkv, seq.d, seq.cmax) q, k, v = _step_inputs(seq, 0, 1) @@ -517,7 +614,14 @@ def export_incache_decode(out_dir: str) -> None: ref_vc = torch.zeros(1, seq.cmax, seq.hkv, seq.d) for t in range(DYN_DECODE_STEPS): cfg = SdpaConfig( - f"incache_{seq.name}_step{t}", seq.hq, seq.hkv, seq.d, 1, seq.cmax, t + f"incache_{seq.name}_step{t}", + seq.hq, + seq.hkv, + seq.d, + 1, + seq.cmax, + t, + kv_f16=seq.kv_f16, ) q, k, v = _step_inputs(seq, t, 1) golden = _golden(cfg, q, k, v, ref_kc, ref_vc).numpy().astype(" torch.Tensor: return x.to(torch.float32) +class ToCopyFloatToIntModule(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x.to(torch.int32) + + +class ToCopyFloatToIntToFloatModule(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x.to(torch.int32).to(torch.float32) + + class ToCopyFloatModule(torch.nn.Module): def forward(self, x: torch.Tensor) -> torch.Tensor: # Same-dtype copy (flat byte-copy path); copy=True keeps the op from @@ -34,11 +46,63 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return x.to(torch.float32, copy=True) -def _export(model: torch.nn.Module, x: torch.Tensor): - ep = torch.export.export(model.eval(), (x,)) - return to_edge_transform_and_lower( - ep, partitioner=[VulkanPartitioner()] - ).to_executorch() +class ToCopyBoolToFloatModule(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x.to(torch.float32) + + +class ToCopyInt8ToFloatModule(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x.to(torch.float32) + + +class CompareToCopyBoolToFloatModule(torch.nn.Module): + def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + return (a > b).to(torch.float32) + + +def to_copy_int_input(shape: tuple[int, ...]) -> torch.Tensor: + n = math.prod(shape) + return (torch.arange(n, dtype=torch.int32) - n // 2).reshape(shape) + + +def to_copy_float_input(shape: tuple[int, ...]) -> torch.Tensor: + n = math.prod(shape) + pattern = torch.tensor( + [-8.75, -3.0, -1.5, -0.25, 0.0, 0.25, 1.5, 3.0, 8.75], + dtype=torch.float32, + ) + repeats = (n + pattern.numel() - 1) // pattern.numel() + return pattern.repeat(repeats)[:n].reshape(shape) + + +def bool_tail_input(shape: tuple[int, ...]) -> torch.Tensor: + n = math.prod(shape) + pattern = torch.tensor([True, False, True, True, False, False, True]) + repeats = (n + pattern.numel() - 1) // pattern.numel() + return pattern.repeat(repeats)[:n].reshape(shape) + + +def compare_to_copy_input_a(shape: tuple[int, ...]) -> torch.Tensor: + n = math.prod(shape) + pattern = torch.tensor([1.0, -1.0, 2.0, -2.0, 3.0, -3.0, 4.0]) + repeats = (n + pattern.numel() - 1) // pattern.numel() + return pattern.repeat(repeats)[:n].reshape(shape) + + +def compare_to_copy_input_b(shape: tuple[int, ...]) -> torch.Tensor: + return torch.zeros(shape, dtype=torch.float32) + + +def _lower(model: torch.nn.Module, *inputs: torch.Tensor): + ep = torch.export.export(model.eval(), inputs) + edge = to_edge_transform_and_lower(ep, partitioner=[VulkanPartitioner()]) + return ep, edge + + +def _export(model: torch.nn.Module, *inputs: torch.Tensor): + _, edge = _lower(model, *inputs) + return edge.to_executorch() def _delegated(et) -> bool: @@ -49,6 +113,29 @@ def _delegated(et) -> bool: ) +def _prepartition_cast_dtypes(ep) -> list[torch.dtype]: + return [ + node.args[1] + for node in ep.graph_module.graph.nodes + if node.op == "call_function" and node.target == torch.ops.aten.to.dtype + ] + + +def _delegated_cast_dtypes(edge) -> list[torch.dtype]: + graph_module = edge.exported_program().graph_module + if any( + "_to_dim_order_copy" in str(getattr(node, "target", "")) + for node in graph_module.graph.nodes + ): + return [] + return [ + node.kwargs["dtype"] + for _, lowered, _ in get_lowered_submodules(graph_module) + for node in lowered.original_module.graph_module.graph.nodes + if "_to_dim_order_copy" in str(getattr(node, "target", "")) + ] + + class ToCopyTest(unittest.TestCase): def test_int_to_float_delegates(self) -> None: x = torch.tensor([1, 2, 3], dtype=torch.int32) @@ -57,9 +144,53 @@ def test_int_to_float_delegates(self) -> None: _delegated(et), "Expected a VulkanBackend delegate (to_copy int->float)" ) + def test_float_to_int_delegates(self) -> None: + x = torch.tensor([-3.75, -1.0, 0.0, 1.9, 63.0], dtype=torch.float32) + et = _export(ToCopyFloatToIntModule(), x) + self.assertTrue( + _delegated(et), "Expected a VulkanBackend delegate (to_copy float->int)" + ) + + def test_roundtrip_keeps_both_casts_in_delegate(self) -> None: + x = torch.tensor([-3.75, -1.0, 0.0, 1.9, 63.0], dtype=torch.float32) + ep, edge = _lower(ToCopyFloatToIntToFloatModule(), x) + expected = [torch.int32, torch.float32] + self.assertEqual(_prepartition_cast_dtypes(ep), expected) + self.assertEqual(_delegated_cast_dtypes(edge), expected) + self.assertTrue(_delegated(edge.to_executorch())) + + for module, one_direction_input in ( + (ToCopyFloatToIntModule(), x), + ( + ToCopyIntToFloatModule(), + torch.tensor([-3, -1, 0, 1, 63], dtype=torch.int32), + ), + ): + _, one_direction_edge = _lower(module, one_direction_input) + self.assertNotEqual(_delegated_cast_dtypes(one_direction_edge), expected) + def test_float_passthrough_delegates(self) -> None: x = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32) et = _export(ToCopyFloatModule(), x) self.assertTrue( _delegated(et), "Expected a VulkanBackend delegate (to_copy float->float)" ) + + def test_bool_to_float_delegates(self) -> None: + x = bool_tail_input((5,)) + ep, edge = _lower(ToCopyBoolToFloatModule(), x) + self.assertEqual(_prepartition_cast_dtypes(ep), [torch.float32]) + self.assertEqual(_delegated_cast_dtypes(edge), [torch.float32]) + self.assertTrue(_delegated(edge.to_executorch())) + + def test_compare_bool_to_float_delegates(self) -> None: + a = compare_to_copy_input_a((5,)) + b = compare_to_copy_input_b((5,)) + ep, edge = _lower(CompareToCopyBoolToFloatModule(), a, b) + self.assertEqual(_prepartition_cast_dtypes(ep), [torch.float32]) + self.assertEqual(_delegated_cast_dtypes(edge), [torch.float32]) + self.assertTrue(_delegated(edge.to_executorch())) + + def test_int8_to_float_does_not_delegate(self) -> None: + x = torch.tensor([-2, 0, 3], dtype=torch.int8) + self.assertFalse(_delegated(_export(ToCopyInt8ToFloatModule(), x))) diff --git a/backends/webgpu/test/test_build_webgpu.sh b/backends/webgpu/test/test_build_webgpu.sh index 1c79d17cf06..ab6b6d022d9 100755 --- a/backends/webgpu/test/test_build_webgpu.sh +++ b/backends/webgpu/test/test_build_webgpu.sh @@ -27,12 +27,14 @@ $PYTHON_EXECUTABLE -m pytest "${SCRIPT_DIR}/test_wgsl_codegen.py" -v echo "=== Step 1: Run Python export tests ===" $PYTHON_EXECUTABLE -m pytest "${SCRIPT_DIR}/ops/test_add.py" -v $PYTHON_EXECUTABLE -m pytest "${SCRIPT_DIR}/ops/test_rms_norm.py" -v +$PYTHON_EXECUTABLE -m pytest "${SCRIPT_DIR}/ops/test_rope_hf.py" -v # ── Step 2: Export .pte model ───────────────────────────────────────────────── echo "=== Step 2: Export test models ===" DISPATCH_ORDER_DIR="/tmp/dispatch_order" PTE_UPDATE_CACHE_MODEL="/tmp/webgpu_update_cache_test.pte" +ROPE_HF_DIR="/tmp/webgpu_rope_hf" cd "${EXECUTORCH_ROOT}" $PYTHON_EXECUTABLE -c " from executorch.backends.webgpu.test.ops.test_dispatch_order import export_dispatch_order_cases @@ -46,6 +48,13 @@ from executorch.backends.webgpu.test.ops.test_update_cache import export_update_ export_update_cache_model('${PTE_UPDATE_CACHE_MODEL}') " || { echo "WARN: update_cache export failed; skipping update_cache native test"; UPDATE_CACHE_OK=0; } +echo "=== Export dynamic HuggingFace RoPE models and goldens ===" +$PYTHON_EXECUTABLE -c " +from executorch.backends.webgpu.test.ops.test_rope_hf import export_rope_hf_dynamic, export_rope_hf_dynamic_sequence +export_rope_hf_dynamic('${ROPE_HF_DIR}') +export_rope_hf_dynamic_sequence('${ROPE_HF_DIR}') +" + echo "=== Export SDPA sweep models (sdpa_.pte + .golden.bin to /tmp) ===" $PYTHON_EXECUTABLE -c " from executorch.backends.webgpu.test.ops.test_sdpa import export_all_sdpa_models @@ -96,6 +105,7 @@ cmake \ "${EXECUTORCH_ROOT}" cmake --build "${NATIVE_BUILD_DIR}" --target webgpu_native_test -j${NPROC} +cmake --build "${NATIVE_BUILD_DIR}" --target webgpu_compute_dispatch_test -j${NPROC} cmake --build "${NATIVE_BUILD_DIR}" --target webgpu_dispatch_order_test -j${NPROC} cmake --build "${NATIVE_BUILD_DIR}" --target webgpu_scratch_buffer_test -j${NPROC} @@ -108,9 +118,11 @@ else fi env \ ${UPDATE_CACHE_ENV_VAR} \ + WEBGPU_TEST_ROPE_HF_DIR="${ROPE_HF_DIR}" \ WEBGPU_TEST_SDPA_DIR=/tmp/ \ "${NATIVE_BUILD_DIR}/backends/webgpu/webgpu_native_test" +"${NATIVE_BUILD_DIR}/backends/webgpu/webgpu_compute_dispatch_test" "${NATIVE_BUILD_DIR}/backends/webgpu/webgpu_dispatch_order_test" "${DISPATCH_ORDER_DIR}" "${NATIVE_BUILD_DIR}/backends/webgpu/webgpu_scratch_buffer_test" diff --git a/backends/webgpu/test/test_cmake_configuration.py b/backends/webgpu/test/test_cmake_configuration.py new file mode 100644 index 00000000000..b42df74b5be --- /dev/null +++ b/backends/webgpu/test/test_cmake_configuration.py @@ -0,0 +1,87 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import pathlib +import re +import unittest + + +def _conditional_branches(source: str, condition: str) -> tuple[str, str]: + lines = source.splitlines() + condition_pattern = re.compile( + rf"^\s*if\s*\(\s*{re.escape(condition)}\s*\)\s*$", + re.IGNORECASE, + ) + command_pattern = re.compile(r"^\s*([A-Za-z_][A-Za-z0-9_]*)\s*\(") + start = next( + (index for index, line in enumerate(lines) if condition_pattern.match(line)), + None, + ) + if start is None: + raise AssertionError(f"if({condition}) branch not found") + + depth = 1 + else_index = None + for index in range(start + 1, len(lines)): + match = command_pattern.match(lines[index]) + if match is None: + continue + command = match.group(1).lower() + if command == "if": + depth += 1 + elif command == "endif": + depth -= 1 + if depth == 0: + if else_index is None: + raise AssertionError(f"if({condition}) has no else() branch") + return ( + "\n".join(lines[start + 1 : else_index]), + "\n".join(lines[else_index + 1 : index]), + ) + elif command == "else" and depth == 1: + if else_index is not None: + raise AssertionError(f"if({condition}) has multiple else() branches") + else_index = index + raise AssertionError(f"if({condition}) has no matching endif()") + + +class TestCMakeConfiguration(unittest.TestCase): + def test_branch_parser_keeps_nested_conditionals_in_native_branch(self) -> None: + source = """ +if ( EMSCRIPTEN ) + wasm_command() +else() + if(APPLE) + apple_command() + else() + linux_command() + endif() +endif() +""" + wasm_branch, native_branch = _conditional_branches(source, "EMSCRIPTEN") + + self.assertIn("wasm_command()", wasm_branch) + self.assertIn("apple_command()", native_branch) + self.assertIn("linux_command()", native_branch) + + def test_emscripten_uses_port_instead_of_native_dawn(self) -> None: + cmake = pathlib.Path(__file__).parents[1] / "CMakeLists.txt" + wasm_branch, native_branch = _conditional_branches( + cmake.read_text(), "EMSCRIPTEN" + ) + port_flag = r'"--use-port=emdawnwebgpu"' + + self.assertRegex( + wasm_branch, + rf"target_compile_options\s*\(\s*webgpu_backend\s+PUBLIC\s+{port_flag}\s*\)", + ) + self.assertRegex( + wasm_branch, + rf"target_link_options\s*\(\s*webgpu_backend\s+PUBLIC\s+{port_flag}\s*\)", + ) + self.assertNotIn("find_package(Dawn", wasm_branch) + self.assertNotIn("--use-port=emdawnwebgpu", native_branch) + self.assertRegex(native_branch, r"find_package\s*\(\s*Dawn\s+REQUIRED\s*\)") diff --git a/backends/webgpu/test/test_native_ci_contract.py b/backends/webgpu/test/test_native_ci_contract.py new file mode 100644 index 00000000000..43b8ec6f56e --- /dev/null +++ b/backends/webgpu/test/test_native_ci_contract.py @@ -0,0 +1,75 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import pathlib +import re +import shlex +import unittest + + +def _bash_array(source: str, name: str) -> list[str]: + match = re.search( + rf"^{re.escape(name)}=\((.*?)\)", source, re.MULTILINE | re.DOTALL + ) + if match is None: + raise AssertionError(f"{name} Bash array not found") + return shlex.split(match.group(1)) + + +class TestNativeCIContract(unittest.TestCase): + def test_builds_and_runs_every_fixed_target(self) -> None: + backend = pathlib.Path(__file__).parents[1] + cmake = (backend / "CMakeLists.txt").read_text() + script = (backend / "scripts/test_webgpu_native_ci.sh").read_text() + required = { + "webgpu_native_test", + "webgpu_dispatch_order_test", + "webgpu_scratch_buffer_test", + "webgpu_update_cache_test", + "webgpu_index_test", + "webgpu_dynamic_shape_test", + "webgpu_dispatch_2d_test", + "webgpu_compute_dispatch_test", + "webgpu_execution_options_test", + "webgpu_output_suppression_test", + "webgpu_op_test_util_test", + } + + self.assertEqual(set(_bash_array(script, "REQUIRED_TARGETS")), required) + self.assertNotIn("not defined in this tree — skipping", script) + self.assertIn('-DPYTHON_EXECUTABLE="${PYTHON_EXECUTABLE}"', script) + self.assertIn("run_with_required_device env WEBGPU_TEST_SDPA_DIR", script) + self.assertIn( + "if ! grep -q '^WebGPU device acquired (native)$' " '<<<"${output}"; then', + script, + ) + for target in required: + self.assertIn(target, cmake) + self.assertIn(f'"${{BIN_DIR}}/{target}"', script) + + def test_requires_symint_and_suppression_fixtures(self) -> None: + script = ( + pathlib.Path(__file__).parents[1] / "scripts/test_webgpu_native_ci.sh" + ).read_text() + + self.assertIn( + "export_output_suppression_models('${OUTPUT_SUPPRESSION_DIR}')", script + ) + self.assertIn('WEBGPU_TEST_SYMINT_BLOB="${SYMINT_BLOB}"', script) + for fixture in ( + "${SYMINT_BLOB}", + "${OUTPUT_SUPPRESSION_DIR}/input.bin", + ): + self.assertIn(f'require_file "{fixture}"', script) + + def test_requires_dynamic_rope_fixture(self) -> None: + script = ( + pathlib.Path(__file__).parents[1] / "scripts/test_webgpu_native_ci.sh" + ).read_text() + + self.assertIn("export_rope_hf_dynamic('${ROPE_HF_DIR}')", script) + self.assertIn('WEBGPU_TEST_ROPE_HF_DIR="${ROPE_HF_DIR}"', script) + self.assertIn('require_file "${ROPE_HF_DIR}/rope_hf_dynamic.pte"', script) diff --git a/backends/webgpu/test/test_webgpu_native.cpp b/backends/webgpu/test/test_webgpu_native.cpp index a21466b9a42..6448568a66e 100644 --- a/backends/webgpu/test/test_webgpu_native.cpp +++ b/backends/webgpu/test/test_webgpu_native.cpp @@ -11,6 +11,7 @@ #include #include #include +#include #include #include #include @@ -20,14 +21,17 @@ #include #include +#include #include #include #include #include #include #include +#include #include #include +#include #include #include @@ -307,13 +311,12 @@ bool sdpa_within_tol( const float* golden, int n, float* ma, - float* mr) { + float* mr, + bool kv_f16 = false) { float atol = 1e-4f, rtol = 1e-3f; - // f16 KV (runtime opt-in) reads K/V at reduced precision; loosen the tol on a - // shader-f16 device to cover that rounding. Harmless for f32 KV (looser - // gate). - const WebGPUContext* kv_ctx = get_default_webgpu_context(); - if (kv_ctx != nullptr && kv_ctx->shader_f16_supported) { + // Only fp16-KV cases receive the tolerance needed for storage rounding; + // device capability alone must not weaken unrelated fp32 tests. + if (kv_f16) { atol = 2e-3f; rtol = 1e-2f; } @@ -512,6 +515,152 @@ std::vector run_rms_norm_at_wg( return result; } +struct RotaryHfProbeParams { + uint32_t n_heads; + uint32_t seq; + uint32_t head_dim; + uint32_t half_dim; + uint32_t num_pairs; + uint32_t rotary_dim; + uint32_t start_pos; + uint32_t _pad; +}; + +std::vector run_rope_hf_2d_probe(const WebGPUContext& ctx) { + constexpr uint32_t kWorkgroupSize = 2; + constexpr uint32_t kWorkgroupsX = 2; + constexpr uint32_t kWorkgroupsY = 2; + constexpr uint32_t kNumPairs = kWorkgroupSize * kWorkgroupsX * kWorkgroupsY; + constexpr uint32_t kHeadDim = kNumPairs * 2; + + std::vector input(kHeadDim); + std::vector output(kHeadDim, 0.0f); + std::vector freqs_cos(kHeadDim, 1.0f); + std::vector freqs_sin(kHeadDim, 0.0f); + for (uint32_t i = 0; i < kHeadDim; i++) { + input[i] = static_cast(i + 1u); + if (i >= kNumPairs) { + freqs_cos[i] = 2.0f; + } + } + + WGPUDevice device = ctx.device; + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kRotaryEmbeddingHfWGSL, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + + WGPUConstantEntry wg_const = {}; + wg_const.key = {"wg_size", WGPU_STRLEN}; + wg_const.value = static_cast(kWorkgroupSize); + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.compute.module = shader; + pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; + pipeline_desc.compute.constantCount = 1; + pipeline_desc.compute.constants = &wg_const; + WGPUComputePipeline pipeline = + wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + WGPUBindGroupLayout layout = + wgpuComputePipelineGetBindGroupLayout(pipeline, 0); + + auto make_buffer = + [device](const void* data, uint64_t size, WGPUBufferUsage usage) { + WGPUBufferDescriptor desc = {}; + desc.size = size; + desc.usage = usage; + desc.mappedAtCreation = true; + WGPUBuffer buffer = wgpuDeviceCreateBuffer(device, &desc); + std::memcpy(wgpuBufferGetMappedRange(buffer, 0, size), data, size); + wgpuBufferUnmap(buffer); + return buffer; + }; + + const uint64_t data_bytes = kHeadDim * sizeof(float); + WGPUBuffer out_buffer = make_buffer( + output.data(), + data_bytes, + WGPUBufferUsage_Storage | WGPUBufferUsage_CopySrc); + WGPUBuffer in_buffer = + make_buffer(input.data(), data_bytes, WGPUBufferUsage_Storage); + WGPUBuffer cos_buffer = + make_buffer(freqs_cos.data(), data_bytes, WGPUBufferUsage_Storage); + WGPUBuffer sin_buffer = + make_buffer(freqs_sin.data(), data_bytes, WGPUBufferUsage_Storage); + const RotaryHfProbeParams params = { + 1u, 1u, kHeadDim, kNumPairs, kNumPairs, kHeadDim, 0u, 0u}; + WGPUBuffer params_buffer = + make_buffer(¶ms, sizeof(params), WGPUBufferUsage_Uniform); + + WGPUBindGroupEntry entries[5] = {}; + const WGPUBuffer buffers[] = { + out_buffer, in_buffer, cos_buffer, sin_buffer, params_buffer}; + const uint64_t sizes[] = { + data_bytes, data_bytes, data_bytes, data_bytes, sizeof(params)}; + for (uint32_t i = 0; i < 5; i++) { + entries[i].binding = i; + entries[i].buffer = buffers[i]; + entries[i].size = sizes[i]; + } + WGPUBindGroupDescriptor bind_group_desc = {}; + bind_group_desc.layout = layout; + bind_group_desc.entryCount = 5; + bind_group_desc.entries = entries; + WGPUBindGroup bind_group = + wgpuDeviceCreateBindGroup(device, &bind_group_desc); + + WGPUBufferDescriptor staging_desc = {}; + staging_desc.size = data_bytes; + staging_desc.usage = WGPUBufferUsage_MapRead | WGPUBufferUsage_CopyDst; + WGPUBuffer staging = wgpuDeviceCreateBuffer(device, &staging_desc); + + WGPUCommandEncoder encoder = wgpuDeviceCreateCommandEncoder(device, nullptr); + WGPUComputePassDescriptor pass_desc = {}; + WGPUComputePassEncoder pass = + wgpuCommandEncoderBeginComputePass(encoder, &pass_desc); + wgpuComputePassEncoderSetPipeline(pass, pipeline); + wgpuComputePassEncoderSetBindGroup(pass, 0, bind_group, 0, nullptr); + wgpuComputePassEncoderDispatchWorkgroups(pass, kWorkgroupsX, kWorkgroupsY, 1); + wgpuComputePassEncoderEnd(pass); + wgpuComputePassEncoderRelease(pass); + wgpuCommandEncoderCopyBufferToBuffer( + encoder, out_buffer, 0, staging, 0, data_bytes); + WGPUCommandBuffer command = wgpuCommandEncoderFinish(encoder, nullptr); + wgpuQueueSubmit(ctx.queue, 1, &command); + wgpuCommandBufferRelease(command); + wgpuCommandEncoderRelease(encoder); + + WgMapData callback = {}; + WGPUBufferMapCallbackInfo callback_info = {}; + callback_info.mode = WGPUCallbackMode_WaitAnyOnly; + callback_info.callback = wg_map_cb; + callback_info.userdata1 = &callback; + WGPUFuture future = wgpuBufferMapAsync( + staging, WGPUMapMode_Read, 0, data_bytes, callback_info); + const WGPUWaitStatus wait = webgpu_wait(ctx.instance, future); + if (wait == WGPUWaitStatus_Success && + callback.status == WGPUMapAsyncStatus_Success) { + const void* mapped = wgpuBufferGetConstMappedRange(staging, 0, data_bytes); + std::memcpy(output.data(), mapped, data_bytes); + wgpuBufferUnmap(staging); + } else { + output.clear(); + } + + wgpuBufferRelease(staging); + wgpuBindGroupRelease(bind_group); + wgpuBufferRelease(params_buffer); + wgpuBufferRelease(sin_buffer); + wgpuBufferRelease(cos_buffer); + wgpuBufferRelease(in_buffer); + wgpuBufferRelease(out_buffer); + wgpuBindGroupLayoutRelease(layout); + wgpuComputePipelineRelease(pipeline); + wgpuShaderModuleRelease(shader); + return output; +} + // linear_q4gsw sweep config; mirrors CONFIGS in test_quantized_linear.py. struct Q4gswConfig { const char* name; @@ -766,6 +915,220 @@ void test_rope( << "apply_rotary_emb exceeds tolerance 1e-3 (abs AND rel)"; } +bool has_shape( + const executorch::aten::Tensor& tensor, + const std::vector& expected) { + if (tensor.dim() != static_cast(expected.size())) { + return false; + } + for (size_t i = 0; i < expected.size(); i++) { + if (tensor.size(static_cast(i)) != expected[i]) { + return false; + } + } + return true; +} + +void test_rope_hf_dynamic(const std::string& dir) { + constexpr int S = 1; + constexpr int NH = 16; + constexpr int NKV = 8; + constexpr int HD = 128; + constexpr int MAXS = 16; + constexpr int positions[] = {0, 7, 15}; + constexpr int xq_numel = S * NH * HD; + constexpr int xk_numel = S * NKV * HD; + constexpr int freqs_numel = MAXS * HD; + + Module module(dir + "rope_hf_dynamic.pte"); + ASSERT_EQ(module.load_forward(), Error::Ok) + << "could not load HF RoPE dynamic model"; + + std::vector xq = load_golden(dir + "rope_hf_dynamic.xq.bin", xq_numel); + std::vector xk = load_golden(dir + "rope_hf_dynamic.xk.bin", xk_numel); + std::vector freqs_cos = + load_golden(dir + "rope_hf_dynamic.freqs_cos.bin", freqs_numel); + std::vector freqs_sin = + load_golden(dir + "rope_hf_dynamic.freqs_sin.bin", freqs_numel); + ASSERT_FALSE( + xq.empty() || xk.empty() || freqs_cos.empty() || freqs_sin.empty()) + << "could not load HF RoPE input binaries from " << dir; + + for (const int position : positions) { + auto xqt = make_tensor_ptr({1, S, NH, HD}, std::vector(xq)); + auto xkt = make_tensor_ptr({1, S, NKV, HD}, std::vector(xk)); + auto fct = make_tensor_ptr({MAXS, HD}, std::vector(freqs_cos)); + auto fst = make_tensor_ptr({MAXS, HD}, std::vector(freqs_sin)); + auto post = make_tensor_ptr( + {1}, std::vector{static_cast(position)}); + auto result = module.forward( + {EValue(xqt), EValue(xkt), EValue(fct), EValue(fst), EValue(post)}); + ASSERT_TRUE(result.ok()) + << "HF RoPE forward failed at position " << position << " (error " + << static_cast(result.error()) << ")"; + const auto& outputs = result.get(); + ASSERT_TRUE( + outputs.size() == 2 && outputs[0].isTensor() && outputs[1].isTensor()) + << "expected exactly two HF RoPE tensor outputs"; + const auto& xq_out = outputs[0].toTensor(); + const auto& xk_out = outputs[1].toTensor(); + ASSERT_TRUE(has_shape(xq_out, {1, S, NH, HD})) + << "HF RoPE query output has the wrong shape at position " << position; + ASSERT_TRUE(has_shape(xk_out, {1, S, NKV, HD})) + << "HF RoPE key output has the wrong shape at position " << position; + + const std::string prefix = + dir + "rope_hf_dynamic.pos" + std::to_string(position); + const std::vector golden_q = + load_golden(prefix + ".xq.golden.bin", xq_numel); + const std::vector golden_k = + load_golden(prefix + ".xk.golden.bin", xk_numel); + ASSERT_FALSE(golden_q.empty() || golden_k.empty()) + << "could not load HF RoPE goldens for position " << position; + + float q_abs = 0.0f, q_rel = 0.0f, k_abs = 0.0f, k_rel = 0.0f; + const bool q_ok = quant_within_tol( + xq_out.const_data_ptr(), + golden_q.data(), + xq_numel, + 1e-4f, + 1e-3f, + &q_abs, + &q_rel); + const bool k_ok = quant_within_tol( + xk_out.const_data_ptr(), + golden_k.data(), + xk_numel, + 1e-4f, + 1e-3f, + &k_abs, + &k_rel); + EXPECT_TRUE(q_ok && k_ok) + << "HF RoPE mismatch at position " << position << ": q abs=" << q_abs + << " rel=" << q_rel << ", k abs=" << k_abs << " rel=" << k_rel; + } + + auto xqt = make_tensor_ptr({1, S, NH, HD}, std::vector(xq)); + auto xkt = make_tensor_ptr({1, S, NKV, HD}, std::vector(xk)); + auto fct = make_tensor_ptr({MAXS, HD}, std::move(freqs_cos)); + auto fst = make_tensor_ptr({MAXS, HD}, std::move(freqs_sin)); + + auto overflow_post = + make_tensor_ptr({1}, std::vector{INT64_C(1) << 32}); + auto overflow = module.forward({ + EValue(xqt), + EValue(xkt), + EValue(fct), + EValue(fst), + EValue(overflow_post), + }); + EXPECT_FALSE(overflow.ok()) + << "HF RoPE accepted a start_pos that aliases to zero when narrowed"; + + auto post = make_tensor_ptr({1}, std::vector{MAXS}); + auto out_of_range = module.forward( + {EValue(xqt), EValue(xkt), EValue(fct), EValue(fst), EValue(post)}); + EXPECT_FALSE(out_of_range.ok()) + << "HF RoPE accepted start_pos + seq beyond the frequency table"; + + auto negative_post = make_tensor_ptr({1}, std::vector{-1}); + auto negative = module.forward({ + EValue(xqt), + EValue(xkt), + EValue(fct), + EValue(fst), + EValue(negative_post), + }); + EXPECT_FALSE(negative.ok()) << "HF RoPE accepted a negative start_pos"; +} + +void test_rope_hf_dynamic_sequence_reused_graph(const std::string& dir) { + constexpr int NH = 16; + constexpr int NKV = 8; + constexpr int HD = 128; + constexpr int MAXS = 16; + struct Case { + int seq; + int position; + }; + constexpr Case cases[] = {{16, 0}, {5, 7}, {1, 15}, {16, 0}}; + + const std::string prefix = dir + "rope_hf_dynamic_sequence"; + Module module(prefix + ".pte"); + ASSERT_EQ(module.load_forward(), Error::Ok) + << "could not load HF RoPE dynamic-sequence model"; + + const int freqs_numel = MAXS * HD; + const std::vector freqs_cos = + load_golden(prefix + ".freqs_cos.bin", freqs_numel); + const std::vector freqs_sin = + load_golden(prefix + ".freqs_sin.bin", freqs_numel); + ASSERT_FALSE(freqs_cos.empty() || freqs_sin.empty()) + << "could not load HF RoPE dynamic-sequence frequencies"; + + for (const Case& c : cases) { + const int xq_numel = c.seq * NH * HD; + const int xk_numel = c.seq * NKV * HD; + const std::string case_prefix = prefix + ".S" + std::to_string(c.seq) + + ".pos" + std::to_string(c.position); + const std::vector xq = + load_golden(case_prefix + ".xq.bin", xq_numel); + const std::vector xk = + load_golden(case_prefix + ".xk.bin", xk_numel); + const std::vector golden_q = + load_golden(case_prefix + ".xq.golden.bin", xq_numel); + const std::vector golden_k = + load_golden(case_prefix + ".xk.golden.bin", xk_numel); + ASSERT_FALSE( + xq.empty() || xk.empty() || golden_q.empty() || golden_k.empty()) + << "could not load HF RoPE dynamic-sequence case " << case_prefix; + + auto xqt = make_tensor_ptr({1, c.seq, NH, HD}, std::vector(xq)); + auto xkt = make_tensor_ptr({1, c.seq, NKV, HD}, std::vector(xk)); + auto fct = make_tensor_ptr({MAXS, HD}, std::vector(freqs_cos)); + auto fst = make_tensor_ptr({MAXS, HD}, std::vector(freqs_sin)); + auto post = make_tensor_ptr( + {1}, std::vector{static_cast(c.position)}); + auto result = module.forward( + {EValue(xqt), EValue(xkt), EValue(fct), EValue(fst), EValue(post)}); + ASSERT_TRUE(result.ok()) + << "HF RoPE dynamic-sequence forward failed for " << case_prefix + << " (error " << static_cast(result.error()) << ")"; + const auto& outputs = result.get(); + ASSERT_TRUE( + outputs.size() == 2 && outputs[0].isTensor() && outputs[1].isTensor()) + << "expected exactly two HF RoPE dynamic-sequence tensor outputs"; + const auto& xq_out = outputs[0].toTensor(); + const auto& xk_out = outputs[1].toTensor(); + ASSERT_TRUE(has_shape(xq_out, {1, c.seq, NH, HD})) + << "HF RoPE query output has the wrong shape for " << case_prefix; + ASSERT_TRUE(has_shape(xk_out, {1, c.seq, NKV, HD})) + << "HF RoPE key output has the wrong shape for " << case_prefix; + + float q_abs = 0.0f, q_rel = 0.0f, k_abs = 0.0f, k_rel = 0.0f; + const bool q_ok = quant_within_tol( + xq_out.const_data_ptr(), + golden_q.data(), + xq_numel, + 1e-4f, + 1e-3f, + &q_abs, + &q_rel); + const bool k_ok = quant_within_tol( + xk_out.const_data_ptr(), + golden_k.data(), + xk_numel, + 1e-4f, + 1e-3f, + &k_abs, + &k_rel); + EXPECT_TRUE(q_ok && k_ok) + << "HF RoPE dynamic-sequence mismatch for " << case_prefix + << ": q abs=" << q_abs << " rel=" << q_rel << ", k abs=" << k_abs + << " rel=" << k_rel; + } +} + void test_prepack( const std::string& model_path, const std::string& golden_path, @@ -888,6 +1251,7 @@ struct SdpaConfig { float denom; // ramp divisor (mirrors Python); small -> large logits bool required = false; // CI (SDPA dir set): absent .pte = FAIL, not skip bool expect_reject = false; // load MUST fail (e.g. D%4 guard), no golden + bool kv_f16 = false; }; const SdpaConfig kSdpaConfigs[] = { @@ -930,6 +1294,28 @@ const SdpaConfig kSdpaConfigs[] = { 0, 16.0f, /*required=*/true}, + {"qwen3_prefill", + 16, + 8, + 128, + 128, + 256, + 0, + 10.0f, + /*required=*/true, + /*expect_reject=*/false, + /*kv_f16=*/true}, + {"qwen3_odd_boundary", + 16, + 8, + 128, + 17, + 64, + 31, + 10.0f, + /*required=*/true, + /*expect_reject=*/false, + /*kv_f16=*/true}, }; // Ramp denominator; mirror of test_sdpa.py::_RAMP_DENOM (keep in sync). @@ -952,9 +1338,8 @@ float sdpa_ramp_t( return static_cast(((i + 31 * t) % mod) - off) / denom; } -// Multi-step replay sequences. Mirror the Python REPLAY_SEQS / Vulkan param -// sets (sdpa_test.cpp:856/867/875). Each seq_lens entry is one step replayed on -// a host-threaded KV cache (big=prefill, mid=multi-token, 1=decode). +// Multi-step replay sequences. The first three mirror Vulkan param sets; Qwen3 +// extends the same Python REPLAY_SEQS contract. struct SdpaSequence { const char* name; int hq; @@ -962,18 +1347,90 @@ struct SdpaSequence { int d; int cmax; std::vector seq_lens; + bool kv_f16 = false; }; const SdpaSequence kSdpaSequences[] = { {"small", 8, 4, 4, 16, {3, 1, 1, 5, 1, 1, 2}}, {"small_d", 6, 2, 8, 16, {3, 1, 1, 5, 1, 1}}, {"llama3", 24, 8, 128, 256, {111, 1, 1, 1, 57, 1, 1}}, + {"qwen3_fd", 16, 8, 128, 64, {17, 1}, /*kv_f16=*/true}, }; +Error load_sdpa_forward(Module& module, bool kv_f16, int sdpa_query_tile = 0) { + if (!kv_f16 && sdpa_query_tile == 0) { + return module.load_forward(); + } + BackendOptions<2> options; + Error error = Error::Ok; + if (kv_f16) { + error = options.set_option("enable_f16_kv_cache", true); + if (error != Error::Ok) { + return error; + } + } + if (sdpa_query_tile != 0) { + error = options.set_option("sdpa_query_tile", sdpa_query_tile); + if (error != Error::Ok) { + return error; + } + } + LoadBackendOptionsMap option_map; + error = option_map.set_options("VulkanBackend", options.view()); + if (error != Error::Ok) { + return error; + } + return module.load_forward(nullptr, nullptr, &option_map); +} + +bool shader_f16_supported_on_test_device() { + const WebGPUContext* context = get_default_webgpu_context(); + return context != nullptr && context->shader_f16_supported; +} + +bool qwen3_q16_supported_on_test_device() { + constexpr uint32_t kQ16StorageBytes = 512u * 4u * sizeof(float) + + 512u * 4u * sizeof(uint16_t) + 128u * 2u * sizeof(float) + + 3u * 16u * sizeof(float); + const WebGPUContext* context = get_default_webgpu_context(); + WGPULimits limits = {}; + return context != nullptr && context->shader_f16_supported && + wgpuDeviceGetLimits(context->device, &limits) == WGPUStatus_Success && + limits.maxComputeWorkgroupSizeX >= 16u && + limits.maxComputeWorkgroupSizeY >= 8u && + limits.maxComputeInvocationsPerWorkgroup >= 128u && + limits.maxComputeWorkgroupStorageSize >= kQ16StorageBytes && + limits.maxStorageBuffersPerShaderStage >= 4u; +} + +bool qwen3_q32_supported_on_test_device() { + constexpr uint32_t kQ32StorageBytes = 1024u * 4u * sizeof(float) + + 512u * 4u * sizeof(uint16_t) + 256u * 2u * sizeof(float) + + 3u * 32u * sizeof(float); + const WebGPUContext* context = get_default_webgpu_context(); + WGPULimits limits = {}; + return context != nullptr && context->shader_f16_supported && + wgpuDeviceGetLimits(context->device, &limits) == WGPUStatus_Success && + limits.maxComputeWorkgroupSizeX >= 32u && + limits.maxComputeWorkgroupSizeY >= 8u && + limits.maxComputeInvocationsPerWorkgroup >= 256u && + limits.maxComputeWorkgroupStorageSize >= kQ32StorageBytes && + limits.maxStorageBuffersPerShaderStage >= 4u; +} + +#ifdef WGPU_BACKEND_ENABLE_PROFILING +constexpr uint32_t kTestRouteMaterializedAttention = 1u << 2; +constexpr uint32_t kTestRouteFlashDecoding = 1u << 10; +constexpr uint32_t kTestRouteK16CausalBound = 1u << 11; +constexpr uint32_t kTestRouteQwen3Q16K16 = 1u << 13; +constexpr uint32_t kTestRouteQwen3Q32K16 = 1u << 14; +#endif // WGPU_BACKEND_ENABLE_PROFILING + void test_sdpa_config( const SdpaConfig& cfg, const std::string& model_path, - const std::string& golden_path) { + const std::string& golden_path, + int sdpa_query_tile = 0) { // Inputs reconstruct test_sdpa.py::_det_inputs bit-for-bit (/16 exact fp32). printf( "\n--- Test: sdpa_with_kv_cache (%s: Hq=%d,Hkv=%d,D=%d,S=%d,Cmax=%d,pos=%d) ---\n", @@ -985,8 +1442,13 @@ void test_sdpa_config( cfg.cmax, cfg.input_pos); + if (cfg.kv_f16 && !shader_f16_supported_on_test_device()) { + printf("SKIP: %s requires shader-f16\n", cfg.name); + return; + } + Module module(model_path); - auto err = module.load_forward(); + auto err = load_sdpa_forward(module, cfg.kv_f16, sdpa_query_tile); if (cfg.expect_reject) { // D not a multiple of 4 must be rejected at load by the head_dim guard. ASSERT_NE(err, Error::Ok) @@ -1031,6 +1493,38 @@ void test_sdpa_config( {EValue(qt), EValue(kt), EValue(vt), EValue(kct), EValue(vct)}); ASSERT_TRUE(result.ok()) << "forward failed (error " << (int)result.error() << ")"; + if (cfg.kv_f16) { +#ifdef WGPU_BACKEND_ENABLE_PROFILING + // Exact Qwen3 geometry + fp16 KV selects the K16 streaming (causal-bound) + // route by default. The sdpa_query_tile RuntimeSpec only swaps the Q16/Q32 + // kernel variant; both map to the K16CausalBound bit. A non-Qwen3 fp16-KV + // shape falls back to the materialized path (or flash-decoding at S==1). + const bool qwen3_geometry = cfg.hq == 16 && cfg.hkv == 8 && cfg.d == 128; + const bool qwen3_streaming = + qwen3_geometry && cfg.s > 1 && qwen3_q16_supported_on_test_device(); + const uint32_t expected_route = qwen3_streaming + ? kTestRouteK16CausalBound + : (cfg.s == 1 ? kTestRouteFlashDecoding + : kTestRouteMaterializedAttention); + EXPECT_EQ( + g_last_route_mask & + (kTestRouteMaterializedAttention | kTestRouteFlashDecoding | + kTestRouteK16CausalBound), + expected_route); + EXPECT_EQ(g_last_route_conflict_count, 0u); + const uint32_t qwen3_tile_routes = + g_last_route_mask & (kTestRouteQwen3Q16K16 | kTestRouteQwen3Q32K16); + if (qwen3_streaming) { + const uint32_t expected_tile_route = + sdpa_query_tile == 32 && qwen3_q32_supported_on_test_device() + ? kTestRouteQwen3Q32K16 + : kTestRouteQwen3Q16K16; + EXPECT_EQ(qwen3_tile_routes, expected_tile_route); + } else { + EXPECT_EQ(qwen3_tile_routes, 0u); + } +#endif // WGPU_BACKEND_ENABLE_PROFILING + } const auto& outputs = result.get(); // Select the attention output [1,S,Hq,D] by shape; the op returns @@ -1062,8 +1556,8 @@ void test_sdpa_config( ASSERT_FALSE(golden.empty()) << "could not load golden " << golden_path; float max_abs_err = 0.0f, max_rel_err = 0.0f; - const bool pass = - sdpa_within_tol(out_data, golden.data(), on, &max_abs_err, &max_rel_err); + const bool pass = sdpa_within_tol( + out_data, golden.data(), on, &max_abs_err, &max_rel_err, cfg.kv_f16); printf( "Max abs error: %e Max rel error: %e (checked %d elements)\n", max_abs_err, @@ -1085,6 +1579,10 @@ void test_sdpa_replay(const SdpaSequence& seq, const std::string& dir) { seq.d, seq.cmax, seq.seq_lens.size()); + if (seq.kv_f16 && !shader_f16_supported_on_test_device()) { + printf("SKIP: %s requires shader-f16\n", seq.name); + return; + } const int cn = seq.cmax * seq.hkv * seq.d; std::vector kc(cn, 0.0f), vc(cn, 0.0f); @@ -1098,7 +1596,7 @@ void test_sdpa_replay(const SdpaSequence& seq, const std::string& dir) { std::to_string(t) + "_S" + std::to_string(s) + "_pos" + std::to_string(input_pos); Module module(base + ".pte"); - ASSERT_EQ(module.load_forward(), Error::Ok) + ASSERT_EQ(load_sdpa_forward(module, seq.kv_f16), Error::Ok) << "could not load " << base << ".pte"; const int qn = s * seq.hq * seq.d; @@ -1124,6 +1622,26 @@ void test_sdpa_replay(const SdpaSequence& seq, const std::string& dir) { {EValue(qt), EValue(kt), EValue(vt), EValue(kct), EValue(vct)}); ASSERT_TRUE(result.ok()) << "forward " << base << ".pte (error " << (int)result.error() << ")"; + if (seq.kv_f16) { +#ifdef WGPU_BACKEND_ENABLE_PROFILING + // S==1 decode -> flash-decoding; a multi-token exact-Qwen3-geometry + // prefill -> the K16 streaming (causal-bound) route by default (no env); + // any other multi-token fp16-KV shape -> materialized. + const bool qwen3_geometry = seq.hq == 16 && seq.hkv == 8 && seq.d == 128; + const bool qwen3_streaming = + qwen3_geometry && qwen3_q16_supported_on_test_device(); + const uint32_t expected_route = s == 1 ? kTestRouteFlashDecoding + : qwen3_streaming ? kTestRouteK16CausalBound + : kTestRouteMaterializedAttention; + EXPECT_EQ( + g_last_route_mask & + (kTestRouteMaterializedAttention | kTestRouteFlashDecoding | + kTestRouteK16CausalBound), + expected_route) + << seq.name << " step" << t; + EXPECT_EQ(g_last_route_conflict_count, 0u) << seq.name << " step" << t; +#endif // WGPU_BACKEND_ENABLE_PROFILING + } const auto& outs = result.get(); // The op returns [k_cache, v_cache, attn_output]: attn has a unique numel; @@ -1171,7 +1689,8 @@ void test_sdpa_replay(const SdpaSequence& seq, const std::string& dir) { ASSERT_FALSE(golden.empty()) << "could not load " << base << ".golden.bin"; const float* ad = outs[attn_idx].toTensor().const_data_ptr(); float ma = 0.0f, mr = 0.0f; - const bool step_ok = sdpa_within_tol(ad, golden.data(), qn, &ma, &mr); + const bool step_ok = + sdpa_within_tol(ad, golden.data(), qn, &ma, &mr, seq.kv_f16); printf( " step%zu (S=%d pos=%d ctx=%d): max abs %e rel %e\n", t, @@ -1462,15 +1981,322 @@ void test_sdpa_incache_decode( } } -// S1 SymInt round-trip: build a graph directly from a dynamic-input_pos SDPA -// blob; confirm input_pos deserializes as a live SymInt and set/read -// round-trips. +void exercise_symint_host_inputs( + WebGPUGraph& graph, + int symint_id, + int input_tensor_id) { + const auto& input_ids = graph.input_ids(); + std::vector inputs(input_ids.size()); + int64_t host_value = 5; + bool found = false; + for (size_t i = 0; i < input_ids.size(); i++) { + if (input_ids[i] == input_tensor_id) { + inputs[i] = {&host_value, sizeof(host_value), true}; + found = true; + } + } + ASSERT_TRUE(found) << "select_as_symint source is not a graph input"; + + const auto update_from_host = [&](int64_t value) { + host_value = value; + graph.update_symints_from_inputs(inputs); + return graph.read_symint(symint_id); + }; + EXPECT_EQ(update_from_host(5), 5); + EXPECT_EQ( + update_from_host(std::numeric_limits::min()), + std::numeric_limits::min()); + EXPECT_EQ( + update_from_host(std::numeric_limits::max()), + std::numeric_limits::max()); + + const auto expect_out_of_range = [&](int64_t value) { + ASSERT_EQ(update_from_host(17), 17); + host_value = value; + try { + graph.update_symints_from_inputs(inputs); + ADD_FAILURE() << "accepted out-of-range select_as_symint value " << value; + } catch (const std::runtime_error& error) { + EXPECT_STREQ( + error.what(), + "select_as_symint: selected value is outside int32 range"); + } + EXPECT_EQ(graph.read_symint(symint_id), 17) + << "rejected value changed the live SymInt"; + }; + expect_out_of_range( + int64_t{std::numeric_limits::min()} - int64_t{1}); + expect_out_of_range( + int64_t{std::numeric_limits::max()} + int64_t{1}); + expect_out_of_range(INT64_C(1) << 32); +} + +void test_symint_input_narrowing() { + namespace vk = vkgraph; + ::flatbuffers::FlatBufferBuilder fbb; + const std::vector dims = {1u}; + std::vector<::flatbuffers::Offset> values; + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::VkTensor, + vk::CreateVkTensorDirect( + fbb, + vk::VkDataType::INT32, + &dims, + /*constant_id=*/-1, + /*mem_obj_id=*/0) + .Union())); + values.push_back(vk::CreateVkValue( + fbb, vk::GraphTypes::Int, vk::CreateInt(fbb, 0).Union())); + values.push_back(vk::CreateVkValue( + fbb, vk::GraphTypes::Int, vk::CreateInt(fbb, 0).Union())); + values.push_back(vk::CreateVkValue( + fbb, vk::GraphTypes::SymInt, vk::CreateSymInt(fbb, 0).Union())); + const std::vector args = {0, 1, 2, 3}; + std::vector<::flatbuffers::Offset> chain; + chain.push_back(vk::CreateOperatorCallDirect( + fbb, 0, "et_vk.select_as_symint.default", &args)); + const std::vector input_ids = {0}; + const std::vector output_ids = {0}; + const auto root = vk::CreateVkGraphDirect( + fbb, "0", &chain, &values, &input_ids, &output_ids); + vk::FinishVkGraphBuffer(fbb, root); + + WebGPUGraph graph; + ASSERT_NO_THROW(graph.build(fbb.GetBufferPointer(), nullptr, 0, nullptr)); + ASSERT_EQ(graph.symint_sources().size(), 1u); + const auto& source = graph.symint_sources().front(); + exercise_symint_host_inputs(graph, source.symint_id, source.input_tensor_id); +} + +void write_u16_le(std::vector& data, size_t offset, uint16_t value) { + data.at(offset) = static_cast(value); + data.at(offset + 1) = static_cast(value >> 8); +} + +void write_u32_le(std::vector& data, size_t offset, uint32_t value) { + for (size_t i = 0; i < sizeof(value); i++) { + data.at(offset + i) = static_cast(value >> (8 * i)); + } +} + +void write_u64_le(std::vector& data, size_t offset, uint64_t value) { + for (size_t i = 0; i < sizeof(value); i++) { + data.at(offset + i) = static_cast(value >> (8 * i)); + } +} + +std::vector make_delegate_header_test_blob() { + std::vector blob(44, 0); + std::memcpy(blob.data() + 4, "VH00", 4); + write_u16_le(blob, 8, 30); + write_u32_le(blob, 10, 32); + write_u32_le(blob, 14, 8); + write_u32_le(blob, 18, 40); + write_u64_le(blob, 22, 4); + return blob; +} + +void finish_inline_constant_graph( + ::flatbuffers::FlatBufferBuilder& fbb, + bool mark_as_kv_cache, + const std::vector& dims, + uint64_t inline_offset = 0) { + namespace vk = vkgraph; + std::vector<::flatbuffers::Offset> values; + const int tensor_count = mark_as_kv_cache ? 5 : 1; + for (int i = 0; i < tensor_count; i++) { + const bool is_cache = mark_as_kv_cache && i >= 3; + const bool is_constant = !mark_as_kv_cache || is_cache; + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::VkTensor, + vk::CreateVkTensorDirect( + fbb, + vk::VkDataType::FLOAT32, + &dims, + is_constant ? (is_cache ? i - 3 : 0) : -1, + is_constant ? -1 : i) + .Union())); + } + + std::vector<::flatbuffers::Offset> chain; + if (mark_as_kv_cache) { + const std::vector args = {0, 1, 2, 3, 4}; + chain.push_back(vk::CreateOperatorCallDirect( + fbb, 0, "sdpa_with_kv_cache.default", &args)); + } + std::vector<::flatbuffers::Offset> constants; + constants.push_back( + vk::CreateVkBytesDirect(fbb, inline_offset, sizeof(float))); + if (mark_as_kv_cache) { + constants.push_back(vk::CreateVkBytesDirect(fbb, 0, sizeof(float))); + } + const std::vector output_ids = {0}; + const auto root = vk::CreateVkGraphDirect( + fbb, "0", &chain, &values, nullptr, &output_ids, &constants); + vk::FinishVkGraphBuffer(fbb, root); +} + +TEST(WebGPUNative, DelegateHeaderRejectsTruncatedRanges) { + const auto blob = make_delegate_header_test_blob(); + EXPECT_TRUE(WebGPUDelegateHeader::parse(blob.data(), blob.size()).ok()); + EXPECT_FALSE(WebGPUDelegateHeader::parse(blob.data(), 29).ok()); + EXPECT_FALSE(WebGPUDelegateHeader::parse(blob.data(), blob.size() - 1).ok()); +} + +TEST(WebGPUNative, InlineConstantExtentIsBounded) { + ::flatbuffers::FlatBufferBuilder fbb; + finish_inline_constant_graph(fbb, false, {1u}); + const std::array data = {0, 0, 0, 0}; + + WebGPUGraph exact_graph; + EXPECT_NO_THROW(exact_graph.build( + fbb.GetBufferPointer(), data.data(), data.size(), nullptr)); + + WebGPUGraph short_graph; + EXPECT_THROW( + short_graph.build( + fbb.GetBufferPointer(), data.data(), data.size() - 1, nullptr), + std::runtime_error); +} + +TEST(WebGPUNative, ZeroByteInlineConstantOffsetIsBounded) { + ::flatbuffers::FlatBufferBuilder fbb; + finish_inline_constant_graph(fbb, false, {0u}, 1); + const std::array data = {0}; + + WebGPUGraph graph; + EXPECT_THROW( + graph.build(fbb.GetBufferPointer(), data.data(), 0, nullptr), + std::runtime_error); +} + +TEST(WebGPUNative, F16KvInlineConstantExtentIsBounded) { + const auto* context = get_default_webgpu_context(); + if (context == nullptr || !context->shader_f16_supported) { + GTEST_SKIP() << "shader-f16 unavailable"; + } + ::flatbuffers::FlatBufferBuilder fbb; + finish_inline_constant_graph(fbb, true, {1u}); + const std::array data = {0, 0, 0, 0}; + + WebGPUGraph graph; + WebGPUGraphConfig config; + config.f16_kv_cache = true; + try { + graph.build( + fbb.GetBufferPointer(), data.data(), data.size() - 1, nullptr, config); + FAIL() << "undersized inline fp16 KV constant was accepted"; + } catch (const std::runtime_error& error) { + EXPECT_STREQ( + error.what(), + "WebGPU f16 KV: inline cache constant exceeds constant data"); + } +} + +void expect_tensor_extent_error( + const std::vector& dims, + const char* expected_error) { + ::flatbuffers::FlatBufferBuilder fbb; + finish_inline_constant_graph(fbb, false, dims); + const std::array data = {0, 0, 0, 0}; + WebGPUGraph graph; + try { + graph.build(fbb.GetBufferPointer(), data.data(), data.size(), nullptr); + ADD_FAILURE() << "overflowing tensor extent was accepted"; + } catch (const std::runtime_error& error) { + EXPECT_STREQ(error.what(), expected_error); + } +} + +TEST(WebGPUNative, TensorExtentOverflowIsRejected) { + expect_tensor_extent_error( + {UINT32_MAX, UINT32_MAX, 2u}, "WebGPU: tensor element count overflows"); + expect_tensor_extent_error( + {UINT32_MAX, UINT32_MAX}, "WebGPU: tensor byte size overflows"); +} + +struct DelegateBlobView { + size_t base_offset; + WebGPUDelegateHeader header; +}; + +std::optional find_delegate_blob( + const std::vector& blob) { + constexpr size_t kHeaderSize = 30; + constexpr size_t kMagicOffset = 4; + constexpr char kMagic[] = {'V', 'H', '0', '0'}; + if (blob.size() < kHeaderSize) { + return std::nullopt; + } + + for (size_t base_offset = 0; base_offset <= blob.size() - kHeaderSize; + base_offset++) { + const uint8_t* base = blob.data() + base_offset; + if (std::memcmp(base + kMagicOffset, kMagic, sizeof(kMagic)) != 0) { + continue; + } + auto header = WebGPUDelegateHeader::parse(base, blob.size() - base_offset); + if (!header.ok()) { + continue; + } + + const uint64_t available = blob.size() - base_offset; + const auto range_is_in_blob = [available](uint64_t offset, uint64_t size) { + return offset <= available && size <= available - offset; + }; + if (!range_is_in_blob(header->flatbuffer_offset, header->flatbuffer_size) || + !range_is_in_blob(header->bytes_offset, header->bytes_size)) { + continue; + } + return DelegateBlobView{base_offset, *header}; + } + return std::nullopt; +} + +TEST(WebGPUNative, StructurallyInvalidVkGraphIsRejectedAtLoad) { + if (g_symint_blob.empty()) { + GTEST_SKIP() << "WEBGPU_TEST_SYMINT_BLOB not set"; + } + FILE* input = std::fopen(g_symint_blob.c_str(), "rb"); + ASSERT_NE(input, nullptr); + std::fseek(input, 0, SEEK_END); + const long file_size = std::ftell(input); + std::fseek(input, 0, SEEK_SET); + ASSERT_GT(file_size, 0); + std::vector blob(static_cast(file_size)); + ASSERT_EQ(std::fread(blob.data(), 1, blob.size(), input), blob.size()); + std::fclose(input); + + const auto delegate = find_delegate_blob(blob); + ASSERT_TRUE(delegate.has_value()); + ASSERT_GE(delegate->header.flatbuffer_size, sizeof(uint32_t)); + const size_t root_offset = + delegate->base_offset + delegate->header.flatbuffer_offset; + std::fill_n(blob.begin() + root_offset, sizeof(uint32_t), UINT8_MAX); + + const std::string malformed_path = "/tmp/webgpu_invalid_vkgraph_" + + std::to_string(reinterpret_cast(blob.data())) + ".pte"; + FILE* output = std::fopen(malformed_path.c_str(), "wb"); + ASSERT_NE(output, nullptr); + ASSERT_EQ(std::fwrite(blob.data(), 1, blob.size(), output), blob.size()); + std::fclose(output); + + Error load_result = Error::Ok; + { + Module module(malformed_path); + load_result = module.load_forward(); + } + EXPECT_NE(load_result, Error::Ok); + EXPECT_EQ(std::remove(malformed_path.c_str()), 0); +} + +// S1 SymInt round-trip: confirm a dynamic input_pos stays live. void test_symint_roundtrip(const std::string& blob_path) { printf("\n--- Test: symint round-trip (%s) ---\n", blob_path.c_str()); FILE* f = std::fopen(blob_path.c_str(), "rb"); - if (!f) { - GTEST_SKIP() << blob_path << " not present"; - } + ASSERT_NE(f, nullptr) << blob_path << " not present"; std::fseek(f, 0, SEEK_END); long n = std::ftell(f); std::fseek(f, 0, SEEK_SET); @@ -1479,13 +2305,17 @@ void test_symint_roundtrip(const std::string& blob_path) { std::fclose(f); ASSERT_EQ(rd, blob.size()) << "short read of " << blob_path; - auto header = WebGPUDelegateHeader::parse(blob.data()); - ASSERT_TRUE(header.ok()) << "delegate header parse"; - const uint8_t* base = blob.data(); + const auto delegate = find_delegate_blob(blob); + ASSERT_TRUE(delegate.has_value()) + << "no complete VH00 delegate blob found in " << blob_path; + const uint8_t* base = blob.data() + delegate->base_offset; WebGPUGraph graph; try { graph.build( - base + header->flatbuffer_offset, base + header->bytes_offset, nullptr); + base + delegate->header.flatbuffer_offset, + base + delegate->header.bytes_offset, + delegate->header.bytes_size, + nullptr); } catch (const std::exception& e) { FAIL() << "graph build: " << e.what(); } @@ -1508,22 +2338,10 @@ void test_symint_roundtrip(const std::string& blob_path) { ASSERT_EQ(graph.read_symint(sid), 7) << "set/read round-trip (got " << graph.read_symint(sid) << ")"; - // Execute-read: feed a fake input_pos=5 via the recorded select_as_symint - // source and confirm update_symints_from_inputs populates the SymInt. const auto& srcs = graph.symint_sources(); ASSERT_FALSE(srcs.empty()) << "no select_as_symint source recorded"; - const auto& in_ids = graph.input_ids(); - std::vector fake_inputs(in_ids.size()); - int64_t fake_pos = 5; - for (size_t i = 0; i < in_ids.size(); i++) { - if (in_ids[i] == srcs[0].input_tensor_id) { - fake_inputs[i] = {&fake_pos, sizeof(int64_t), true}; - } - } - graph.update_symints_from_inputs(fake_inputs); - ASSERT_EQ(graph.read_symint(srcs[0].symint_id), 5) - << "execute-read (got " << graph.read_symint(srcs[0].symint_id) - << ", want 5)"; + exercise_symint_host_inputs( + graph, srcs[0].symint_id, srcs[0].input_tensor_id); printf( "PASS: symint round-trip (SymInt %d: deserialize, live buffer, " @@ -1537,9 +2355,7 @@ void test_symint_roundtrip(const std::string& blob_path) { void test_resize_hook(const std::string& blob_path) { printf("\n--- Test: resize-hook dirty-gating (%s) ---\n", blob_path.c_str()); FILE* f = std::fopen(blob_path.c_str(), "rb"); - if (!f) { - GTEST_SKIP() << blob_path << " not present"; - } + ASSERT_NE(f, nullptr) << blob_path << " not present"; std::fseek(f, 0, SEEK_END); long n = std::ftell(f); std::fseek(f, 0, SEEK_SET); @@ -1547,13 +2363,17 @@ void test_resize_hook(const std::string& blob_path) { size_t rd = std::fread(blob.data(), 1, blob.size(), f); std::fclose(f); ASSERT_EQ(rd, blob.size()) << "short read of " << blob_path; - auto header = WebGPUDelegateHeader::parse(blob.data()); - ASSERT_TRUE(header.ok()) << "delegate header parse"; - const uint8_t* base = blob.data(); + const auto delegate = find_delegate_blob(blob); + ASSERT_TRUE(delegate.has_value()) + << "no complete VH00 delegate blob found in " << blob_path; + const uint8_t* base = blob.data() + delegate->base_offset; WebGPUGraph graph; try { graph.build( - base + header->flatbuffer_offset, base + header->bytes_offset, nullptr); + base + delegate->header.flatbuffer_offset, + base + delegate->header.bytes_offset, + delegate->header.bytes_size, + nullptr); } catch (const std::exception& e) { FAIL() << "graph build: " << e.what(); } @@ -1702,7 +2522,7 @@ static bool test_slice_double_start_case(double start_d, int out_len) { WebGPUGraph graph; try { - graph.build(fbb.GetBufferPointer(), nullptr, nullptr); + graph.build(fbb.GetBufferPointer(), nullptr, 0, nullptr); } catch (const std::exception& e) { printf("FAIL: graph build threw: %s\n", e.what()); return false; @@ -1715,12 +2535,13 @@ static bool test_slice_double_start_case(double start_d, int out_len) { std::vector inputs(1); inputs[0] = {in.data(), in.size() * sizeof(float), false}; std::vector out(out_len, -1.0f); - std::vector> outputs(1); - outputs[0] = {out.data(), out.size() * sizeof(float)}; + std::vector outputs(1); + outputs[0] = {out.data(), out.size() * sizeof(float), true}; try { graph.copy_inputs(inputs); - graph.execute(); - graph.copy_outputs(outputs); + const WebGPUExecutionPlan plan = graph.make_execution_plan({}); + graph.execute(plan); + graph.copy_outputs(outputs, plan); } catch (const std::exception& e) { printf("FAIL: slice execute threw: %s\n", e.what()); return false; @@ -1799,7 +2620,7 @@ static bool test_slice_double_start_rejects(double bad_start) { WebGPUGraph graph; try { - graph.build(fbb.GetBufferPointer(), nullptr, nullptr); + graph.build(fbb.GetBufferPointer(), nullptr, 0, nullptr); } catch (const std::exception& e) { printf("PASS: rejected as expected: %s\n", e.what()); return true; @@ -1897,7 +2718,7 @@ static bool test_select_double_scalar_case( WebGPUGraph graph; try { - graph.build(fbb.GetBufferPointer(), nullptr, nullptr); + graph.build(fbb.GetBufferPointer(), nullptr, 0, nullptr); } catch (const std::exception& e) { printf("FAIL: graph build threw: %s\n", e.what()); return false; @@ -1907,12 +2728,13 @@ static bool test_select_double_scalar_case( std::vector inputs = { {in.data(), in.size() * sizeof(float), false}}; std::vector out(expected.size(), -1.0f); - std::vector> outputs = { - {out.data(), out.size() * sizeof(float)}}; + std::vector outputs = { + {out.data(), out.size() * sizeof(float), true}}; try { graph.copy_inputs(inputs); - graph.execute(); - graph.copy_outputs(outputs); + const WebGPUExecutionPlan plan = graph.make_execution_plan({}); + graph.execute(plan); + graph.copy_outputs(outputs, plan); } catch (const std::exception& e) { printf("FAIL: select execute threw: %s\n", e.what()); return false; @@ -1936,7 +2758,7 @@ static bool test_select_scalar_build_error( WebGPUGraph graph; try { - graph.build(fbb.GetBufferPointer(), nullptr, nullptr); + graph.build(fbb.GetBufferPointer(), nullptr, 0, nullptr); } catch (const std::exception& e) { const std::string error = e.what(); if (error.find(expected_error) != std::string::npos) { @@ -2008,6 +2830,82 @@ static bool test_select_double_scalars() { return ok; } +void expect_rope_hf_resize_numel_overflow(uint32_t q_heads, uint32_t k_heads) { + namespace vk = vkgraph; + ::flatbuffers::FlatBufferBuilder fbb; + + const std::vector q_dims = {1u, 1u, q_heads, 2u}; + const std::vector k_dims = {1u, 1u, k_heads, 2u}; + const std::vector freqs_dims = {2u, 2u}; + std::vector<::flatbuffers::Offset> values; + const auto add_tensor = [&](const std::vector& dims, int mem_id) { + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::VkTensor, + vk::CreateVkTensorDirect( + fbb, + vk::VkDataType::FLOAT32, + &dims, + /*constant_id=*/-1, + /*mem_obj_id=*/mem_id) + .Union())); + }; + add_tensor(q_dims, 0); + add_tensor(k_dims, 1); + add_tensor(freqs_dims, 2); + add_tensor(freqs_dims, 3); + values.push_back(vk::CreateVkValue( + fbb, vk::GraphTypes::Int, vk::CreateInt(fbb, 0).Union())); + add_tensor(q_dims, 4); + add_tensor(k_dims, 5); + const std::vector output_items = {5, 6}; + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::ValueList, + vk::CreateValueListDirect(fbb, &output_items).Union())); + + const std::vector args = {0, 1, 2, 3, 4, 7}; + std::vector<::flatbuffers::Offset> chain; + chain.push_back(vk::CreateOperatorCallDirect( + fbb, 0, "et_vk.apply_rotary_emb_hf.default", &args)); + const std::vector input_ids = {0, 1, 2, 3}; + const std::vector output_ids = {5, 6}; + const auto root = vk::CreateVkGraphDirect( + fbb, "0", &chain, &values, &input_ids, &output_ids); + vk::FinishVkGraphBuffer(fbb, root); + + WebGPUGraph graph; + ASSERT_NO_THROW(graph.build(fbb.GetBufferPointer(), nullptr, 0, nullptr)); + ASSERT_EQ(graph.num_dispatches(), 2u); + const uint32_t q_x = graph.dispatch_at(0).workgroup_count_x; + const uint32_t q_y = graph.dispatch_at(0).workgroup_count_y; + const uint32_t k_x = graph.dispatch_at(1).workgroup_count_x; + const uint32_t k_y = graph.dispatch_at(1).workgroup_count_y; + + constexpr int64_t kLargeBatch = INT64_C(1) << 30; + const std::vector q_live = { + kLargeBatch, 1, static_cast(q_heads), 2}; + const std::vector k_live = { + kLargeBatch, 1, static_cast(k_heads), 2}; + graph.get_tensor(0).dims = q_live; + graph.get_tensor(1).dims = k_live; + ASSERT_NO_THROW(graph.resize_input(0, q_live)); + ASSERT_NO_THROW(graph.resize_input(1, k_live)); + + try { + graph.propagate_resize(); + FAIL() << "accepted q/k element count outside uint32 range"; + } catch (const std::runtime_error& error) { + EXPECT_STREQ( + error.what(), + "apply_rotary_emb_hf(resize): element index exceeds uint32 range"); + } + EXPECT_EQ(graph.dispatch_at(0).workgroup_count_x, q_x); + EXPECT_EQ(graph.dispatch_at(0).workgroup_count_y, q_y); + EXPECT_EQ(graph.dispatch_at(1).workgroup_count_x, k_x); + EXPECT_EQ(graph.dispatch_at(1).workgroup_count_y, k_y); +} + // apply_rotary_emb on-GPU configs: multi + decode (env-gated, run-if-present). struct RopeConfig { const char* name; @@ -2181,6 +3079,48 @@ TEST(WebGPUNative, Rope) { } } +TEST(WebGPUNative, RopeHfDynamic) { + const char* env = std::getenv("WEBGPU_TEST_ROPE_HF_DIR"); + if (env == nullptr || *env == '\0') { + GTEST_SKIP() << "WEBGPU_TEST_ROPE_HF_DIR not set"; + } + std::string dir = env; + if (dir.back() != '/') { + dir += '/'; + } + test_rope_hf_dynamic(dir); +} + +TEST(WebGPUNative, RopeHfDynamicSequenceReusedGraph) { + const char* env = std::getenv("WEBGPU_TEST_ROPE_HF_DIR"); + if (env == nullptr || *env == '\0') { + GTEST_SKIP() << "WEBGPU_TEST_ROPE_HF_DIR not set"; + } + std::string dir = env; + if (dir.back() != '/') { + dir += '/'; + } + test_rope_hf_dynamic_sequence_reused_graph(dir); +} + +TEST(WebGPUNative, RopeHfUsesFull2DGridStride) { + const WebGPUContext* ctx = get_default_webgpu_context(); + ASSERT_NE(ctx, nullptr); + const std::vector output = run_rope_hf_2d_probe(*ctx); + ASSERT_EQ(output.size(), 16u) << "HF RoPE probe output map failed"; + for (size_t i = 0; i < output.size(); i++) { + const float expected = + static_cast(i + 1u) * (i < output.size() / 2u ? 1.0f : 2.0f); + EXPECT_EQ(output[i], expected) + << "HF RoPE 2D grid or second-half frequency mismatch at element " << i; + } +} + +TEST(WebGPUNative, RopeHfResizeRejectsQOrKNumelOverflow) { + expect_rope_hf_resize_numel_overflow(/*q_heads=*/2, /*k_heads=*/1); + expect_rope_hf_resize_numel_overflow(/*q_heads=*/1, /*k_heads=*/2); +} + TEST(WebGPUNative, Prepack) { if (g_prepack_model_path.empty() || g_prepack_golden_path.empty()) { GTEST_SKIP() << "WEBGPU_TEST_PREPACK_MODEL/GOLDEN not set"; @@ -2216,6 +3156,95 @@ TEST(WebGPUNative, PrepackTied) { } // SDPA sweep: configs self-discover sdpa_.pte; required=FAIL else skip. +TEST(WebGPUNative, Qwen3SdpaFixtureContract) { + const auto find_config = [](const char* name) { + return std::find_if( + std::begin(kSdpaConfigs), + std::end(kSdpaConfigs), + [name](const SdpaConfig& cfg) { + return std::strcmp(cfg.name, name) == 0; + }); + }; + const auto prefill = find_config("qwen3_prefill"); + const auto boundary = find_config("qwen3_odd_boundary"); + ASSERT_NE(prefill, std::end(kSdpaConfigs)); + ASSERT_NE(boundary, std::end(kSdpaConfigs)); + EXPECT_EQ( + std::vector( + {prefill->hq, + prefill->hkv, + prefill->d, + prefill->s, + prefill->cmax, + prefill->input_pos}), + std::vector({16, 8, 128, 128, 256, 0})); + EXPECT_EQ( + std::vector( + {boundary->hq, + boundary->hkv, + boundary->d, + boundary->s, + boundary->cmax, + boundary->input_pos}), + std::vector({16, 8, 128, 17, 64, 31})); + EXPECT_TRUE(prefill->kv_f16 && boundary->kv_f16); + + const auto replay = std::find_if( + std::begin(kSdpaSequences), + std::end(kSdpaSequences), + [](const SdpaSequence& seq) { + return std::strcmp(seq.name, "qwen3_fd") == 0; + }); + ASSERT_NE(replay, std::end(kSdpaSequences)); + EXPECT_EQ( + std::vector({replay->hq, replay->hkv, replay->d, replay->cmax}), + std::vector({16, 8, 128, 64})); + EXPECT_EQ(replay->seq_lens, std::vector({17, 1})); + EXPECT_TRUE(replay->kv_f16); +} + +TEST(WebGPUNative, Qwen3SdpaRoutes) { + if (g_sdpa_dir.empty()) { + GTEST_SKIP() << "WEBGPU_TEST_SDPA_DIR not set"; + } + if (!qwen3_q16_supported_on_test_device()) { + GTEST_SKIP() << "Qwen3 Q16 K16 device limits unavailable"; + } + // Default route: exact-Qwen3-geometry fp16-KV configs select the Q16 K16 + // streaming (causal-bound) route by geometry (no runtime config needed) -- + // the per-config assertions live in test_sdpa_config / test_sdpa_replay. + for (const auto& cfg : kSdpaConfigs) { + if (std::strncmp(cfg.name, "qwen3_", 6) != 0) { + continue; + } + const std::string base = g_sdpa_dir + "sdpa_" + cfg.name; + test_sdpa_config(cfg, base + ".pte", base + ".golden.bin"); + } + const auto replay = std::find_if( + std::begin(kSdpaSequences), + std::end(kSdpaSequences), + [](const SdpaSequence& seq) { + return std::strcmp(seq.name, "qwen3_fd") == 0; + }); + ASSERT_NE(replay, std::end(kSdpaSequences)); + test_sdpa_replay(*replay, g_sdpa_dir); + + // Run Q32 over both an aligned prefill and the S=17/nonzero-position case so + // the partial final workgroup's row mask is covered. Unsupported Q32 devices + // intentionally fall back to the already-qualified Q16 route. + for (const auto& cfg : kSdpaConfigs) { + if (std::strncmp(cfg.name, "qwen3_", 6) != 0) { + continue; + } + const std::string base = g_sdpa_dir + "sdpa_" + cfg.name; + test_sdpa_config( + cfg, + base + ".pte", + base + ".golden.bin", + /*sdpa_query_tile=*/32); + } +} + TEST(WebGPUNative, SdpaSweep) { const std::string& dir = g_sdpa_dir; bool ran = false; @@ -2345,7 +3374,8 @@ TEST(WebGPUNative, SdpaAllFamiliesRanWhenDirSet) { TEST(WebGPUNative, SymintRoundtrip) { if (g_symint_blob.empty()) { - GTEST_SKIP() << "WEBGPU_TEST_SYMINT_BLOB not set"; + test_symint_input_narrowing(); + return; } test_symint_roundtrip(g_symint_blob); } diff --git a/backends/webgpu/test/test_wgsl_codegen.py b/backends/webgpu/test/test_wgsl_codegen.py index 46d285aa60b..d293867e7ab 100644 --- a/backends/webgpu/test/test_wgsl_codegen.py +++ b/backends/webgpu/test/test_wgsl_codegen.py @@ -9,12 +9,17 @@ Loads the generator by file path (no package/namespace dependency). """ +import contextlib import hashlib import importlib.util +import io +import os import re +import stat import tempfile import unittest from pathlib import Path +from unittest import mock import yaml @@ -72,6 +77,39 @@ def _function_source(text: str, name: str) -> str: class WgslCodegenTest(unittest.TestCase): + def test_compare_word_count_does_not_overflow_u32(self) -> None: + source = (g.BACKEND_ROOT / "runtime/ops/compare/compare.wgsl").read_text() + expression = "(params.num_elements - 1u) / 4u + 1u" + self.assertIn(expression, source) + for num_elements in (1, 4, 5, (1 << 32) - 3, (1 << 32) - 2, (1 << 32) - 1): + self.assertEqual( + (num_elements - 1) // 4 + 1, + (num_elements + 3) // 4, + ) + + def test_registry_entries_match_concrete_headers(self) -> None: + entries = g.registry_entries() + names = [entry.name for entry in entries] + expected = sorted( + header.name[: -len("_wgsl.h")] + for wgsl in g.discover() + for header, _ in g.headers_for_shader(wgsl) + ) + self.assertEqual(names, expected) + self.assertEqual(len(names), len(set(names))) + + def test_registry_render_is_deterministic(self) -> None: + entries = g.registry_entries() + self.assertEqual( + g.render_registry(entries), + g.render_registry(list(reversed(entries))), + ) + + def test_registry_rejects_duplicate_names(self) -> None: + entry = g.registry_entries()[0] + with self.assertRaisesRegex(ValueError, "duplicate shader registry name"): + g.render_registry([entry, entry]) + def test_symbol_base(self) -> None: self.assertEqual(g.symbol_base("binary_add"), "BinaryAdd") self.assertEqual( @@ -133,6 +171,22 @@ def test_render_header_embeds_sha256(self) -> None: self.assertEqual(g.embedded_sha256(h), want) self.assertEqual(g.wgsl_sha256(wgsl), want) + def test_render_header_long_name_is_clang_format_stable(self) -> None: + stem = "streaming_attention_qwen3_q32_k16_causal_bound" + wgsl = "@compute @workgroup_size(32, 8, 1)\nfn main(){}\n" + h = g.render_header(Path(f"runtime/ops/sdpa/{stem}.wgsl"), wgsl) + + self.assertIn( + f"// @generated from {stem}.wgsl\n// DO NOT EDIT.", + h, + ) + self.assertIn( + "inline constexpr uint32_t\n" + " kStreamingAttentionQwen3Q32K16CausalBoundWorkgroupSizeX = 32;", + h, + ) + self.assertEqual(g.embedded_sha256(h), g.wgsl_sha256(wgsl)) + def test_embedded_sha256_missing_returns_empty(self) -> None: self.assertEqual(g.embedded_sha256("no sha line here\n"), "") @@ -153,6 +207,78 @@ def test_committed_headers_match_generator(self) -> None: got, want, f"{header.name} stale; run scripts/gen_wgsl_headers.py" ) + def test_generated_output_manifest_digest(self) -> None: + outputs = sorted( + [ + *(g.BACKEND_ROOT / "runtime/ops").glob("**/*_wgsl.h"), + g.registry_path(), + ] + ) + digest = hashlib.sha256() + for output in outputs: + digest.update(output.relative_to(g.BACKEND_ROOT).as_posix().encode()) + digest.update(b"\0") + digest.update(output.read_bytes()) + digest.update(b"\0") + self.assertEqual(len(outputs), 135) + self.assertEqual( + digest.hexdigest(), + "61bcf7671ba11a1d9f26e6488d7c62de855728e6cef6ef7983e2a04c4e1c79ca", + ) + self.assertEqual( + hashlib.sha256(g.registry_path().read_bytes()).hexdigest(), + "b3e84b54e6fcb7e48477018d91de62977202418f3b45cb6c28a70dfe86944176", + ) + + def test_rope_hf_reconstructs_full_2d_grid_stride(self) -> None: + shader = ( + g.BACKEND_ROOT / "runtime" / "ops" / "rope" / "rotary_embedding_hf.wgsl" + ).read_text() + self.assertIn("@builtin(num_workgroups) num_workgroups", shader) + self.assertIn( + "gid.x + gid.y * (num_workgroups.x * wg_size)", + shader, + ) + self.assertIn("let freqs_b_idx = freqs_a_idx + half_dim;", shader) + self.assertIn("t_out[b_idx] = x_b * c_b + x_a * si_b;", shader) + + wg_size = 2 + workgroups_x = 2 + indices = [ + group_x * wg_size + lane + group_y * (workgroups_x * wg_size) + for group_y in range(2) + for group_x in range(workgroups_x) + for lane in range(wg_size) + ] + self.assertEqual(indices, list(range(8))) + + def test_qwen3_runtime_eligibility_is_exact(self) -> None: + sdpa = (g.BACKEND_ROOT / "runtime/ops/sdpa/Sdpa.cpp").read_text() + self.assertIn("q/k/v/output must be fp32", sdpa) + self.assertIn("cache dtype does not match the selected storage mode", sdpa) + self.assertIn("scale == qwen3_expected_scale", sdpa) + self.assertNotIn("std::fabs(scale - qwen3_expected_scale)", sdpa) + + def test_fp16_kv_graph_guards_transfer_and_topology(self) -> None: + graph = (g.BACKEND_ROOT / "runtime/WebGPUGraph.cpp").read_text() + self.assertIn("serialized cache tensor must be fp32", graph) + self.assertIn("consumed through a ValueList", graph) + self.assertIn("preserve it while changing storage", graph) + + copy_inputs = graph.index("void WebGPUGraph::copy_inputs") + input_guard = graph.index( + "fp16 device input requires an fp32 host tensor", copy_inputs + ) + fast_path = graph.index("// Fast path", copy_inputs) + self.assertLess(input_guard, fast_path) + + copy_outputs = graph.index("void WebGPUGraph::copy_outputs") + output_guard = graph.index( + "fp16 device output requires an fp32 host tensor", copy_outputs + ) + map_request = graph.index("wgpuBufferMapAsync", copy_outputs) + self.assertLess(output_guard, map_request) + def test_parse_workgroup_allows_space(self) -> None: # @workgroup_size (64) — the spec-legal spaced form must still parse. self.assertEqual( @@ -179,7 +305,12 @@ def test_check_fails_on_stale_header(self) -> None: orig = g.BACKEND_ROOT g.BACKEND_ROOT = Path(tmp) try: - self.assertEqual(g.main(["--check"]), 1) + output = io.StringIO() + with contextlib.redirect_stdout(output): + self.assertEqual(g.main(["--check"]), 1) + self.assertEqual( + output.getvalue().count("Stale embedded WGSL headers"), 1 + ) finally: g.BACKEND_ROOT = orig @@ -239,6 +370,352 @@ def test_render_header_3d_emits_xyz(self) -> None: self.assertIn("inline constexpr uint32_t kFooWorkgroupSizeZ = 2;", h) +class WgslGenerationTransactionTest(unittest.TestCase): + _VALID_SHADER = "@compute @workgroup_size(1)\nfn main() {}\n" + + def setUp(self) -> None: + self._tmp = tempfile.TemporaryDirectory() + self.root = Path(self._tmp.name) + (self.root / "runtime/ops").mkdir(parents=True) + self._original_root = g.BACKEND_ROOT + g.BACKEND_ROOT = self.root + + def tearDown(self) -> None: + g.BACKEND_ROOT = self._original_root + self._tmp.cleanup() + + def _write_shader( + self, directory: str, stem: str, text: str = _VALID_SHADER + ) -> Path: + op_dir = self.root / "runtime/ops" / directory + op_dir.mkdir(parents=True, exist_ok=True) + shader = op_dir / f"{stem}.wgsl" + shader.write_text(text) + return shader + + def _write_template( + self, directory: str, stem: str, text: str, names: list[str] + ) -> Path: + shader = self._write_shader(directory, stem, text) + spec = { + stem: { + "parameter_names_with_default_values": {}, + "shader_variants": [{"NAME": name} for name in names], + } + } + shader.with_suffix(".yaml").write_text(yaml.safe_dump(spec)) + return shader + + def _snapshot(self): + return { + path.relative_to(self.root).as_posix(): ( + path.read_bytes(), + stat.S_IMODE(path.stat().st_mode), + ) + for path in sorted(self.root.rglob("*")) + if path.is_file() + } + + def _run(self, *args: str): + output = io.StringIO() + with contextlib.redirect_stdout(output): + result = g.main(list(args)) + return result, output.getvalue() + + def _assert_no_temps(self) -> None: + self.assertEqual(list(self.root.rglob("*.tmp")), []) + + @staticmethod + def _fail_nth(real_fn, n: int): + calls = 0 + + def wrapped(*args, **kwargs): + nonlocal calls + calls += 1 + if calls == n: + raise OSError(f"injected failure on call {n}") + return real_fn(*args, **kwargs) + + return wrapped + + def test_late_malformed_shader_leaves_tree_unchanged(self) -> None: + good = self._write_shader("a", "good") + good.with_name("good_wgsl.h").write_text("stale\n") + self._write_shader("z", "bad", "${MISSING\n") + before = self._snapshot() + + result, _ = self._run() + + self.assertEqual(result, 1) + self.assertEqual(self._snapshot(), before) + + def test_duplicate_registry_name_leaves_tree_unchanged(self) -> None: + self._write_shader("a", "shared") + self._write_shader("b", "shared") + before = self._snapshot() + + result, _ = self._run() + + self.assertEqual(result, 1) + self.assertEqual(self._snapshot(), before) + + def test_duplicate_registry_symbol_leaves_tree_unchanged(self) -> None: + self._write_shader("a", "foo_bar") + self._write_shader("b", "foo__bar") + before = self._snapshot() + + result, _ = self._run() + + self.assertEqual(result, 1) + self.assertEqual(self._snapshot(), before) + + def test_duplicate_output_path_leaves_tree_unchanged(self) -> None: + self._write_template("op", "op", self._VALID_SHADER, ["duplicate", "duplicate"]) + before = self._snapshot() + + result, _ = self._run() + + self.assertEqual(result, 1) + self.assertEqual(self._snapshot(), before) + + def test_second_stage_failure_leaves_tree_unchanged(self) -> None: + self._write_shader("a", "first") + self._write_shader("b", "second") + before = self._snapshot() + real_mkstemp = tempfile.mkstemp + + with mock.patch( + "tempfile.mkstemp", side_effect=self._fail_nth(real_mkstemp, 2) + ): + result, _ = self._run() + + self.assertEqual(result, 1) + self.assertEqual(self._snapshot(), before) + self._assert_no_temps() + + def test_staging_interrupt_leaves_tree_unchanged(self) -> None: + self._write_shader("a", "first") + self._write_shader("b", "second") + before = self._snapshot() + real_chmod = Path.chmod + + def interrupt_second(path, mode, **kwargs): + interrupt_second.calls += 1 + if interrupt_second.calls == 2: + raise KeyboardInterrupt("injected staging interruption") + return real_chmod(path, mode, **kwargs) + + interrupt_second.calls = 0 + with mock.patch.object( + Path, "chmod", autospec=True, side_effect=interrupt_second + ): + with self.assertRaises(KeyboardInterrupt): + self._run() + + self.assertEqual(self._snapshot(), before) + self._assert_no_temps() + + def test_replace_failure_restores_existing_destination(self) -> None: + self._write_shader("op", "op") + registry = g.registry_path() + registry.write_text("old registry\n") + registry.chmod(0o600) + before = self._snapshot() + real_replace = os.replace + + with mock.patch("os.replace", side_effect=self._fail_nth(real_replace, 2)): + result, _ = self._run() + + self.assertEqual(result, 1) + self.assertEqual(self._snapshot(), before) + self._assert_no_temps() + + def test_replace_failure_removes_new_destination(self) -> None: + self._write_shader("op", "op") + before = self._snapshot() + real_replace = os.replace + + with mock.patch("os.replace", side_effect=self._fail_nth(real_replace, 2)): + result, _ = self._run() + + self.assertEqual(result, 1) + self.assertEqual(self._snapshot(), before) + self._assert_no_temps() + + def test_multiple_rollback_errors_do_not_stop_later_restores(self) -> None: + headers = [] + for directory in ("a", "b", "c"): + shader = self._write_shader(directory, directory) + header = shader.with_name(f"{directory}_wgsl.h") + header.write_text(f"old {directory}\n") + headers.append(header) + registry = g.registry_path() + registry.write_text("old registry\n") + real_replace = os.replace + calls = [] + + def fail_commit_and_two_rollbacks(source, destination): + calls.append(Path(destination)) + if len(calls) in (4, 5, 6): + raise OSError(f"injected failure on replace {len(calls)}") + return real_replace(source, destination) + + with mock.patch("os.replace", side_effect=fail_commit_and_two_rollbacks): + result, output = self._run() + + self.assertEqual(result, 1) + self.assertEqual(len(calls), 7) + self.assertEqual(calls[-3:], [headers[1], headers[0], registry]) + self.assertIn(f"cannot roll back {headers[1]}", output) + self.assertIn(f"cannot roll back {headers[0]}", output) + self.assertEqual(registry.read_text(), "old registry\n") + self.assertNotEqual(headers[0].read_text(), "old a\n") + self.assertNotEqual(headers[1].read_text(), "old b\n") + self.assertEqual(headers[2].read_text(), "old c\n") + self._assert_no_temps() + + def test_success_preserves_existing_mode_and_creates_0644(self) -> None: + shader = self._write_shader("op", "op") + registry = g.registry_path() + registry.write_text("old registry\n") + registry.chmod(0o600) + + result, _ = self._run() + + self.assertEqual(result, 0) + self.assertEqual(stat.S_IMODE(registry.stat().st_mode), 0o600) + self.assertEqual( + stat.S_IMODE(shader.with_name("op_wgsl.h").stat().st_mode), 0o644 + ) + + def test_orphans_are_sorted_reported_and_never_deleted(self) -> None: + self._write_shader("new", "new") + orphan_z = self.root / "runtime/ops/z/old_z_wgsl.h" + orphan_a = self.root / "runtime/ops/a/old_a_wgsl.h" + orphan_z.parent.mkdir(parents=True) + orphan_a.parent.mkdir(parents=True) + orphan_z.write_text("// @generated\n") + orphan_a.write_text("// @generated\n") + before = self._snapshot() + + check_result, check_output = self._run("--check") + normal_result, normal_output = self._run() + + self.assertEqual(check_result, 1) + self.assertEqual(normal_result, 1) + for output in (check_output, normal_output): + self.assertIn("Orphan", output) + self.assertLess(output.index("old_a_wgsl.h"), output.index("old_z_wgsl.h")) + self.assertEqual(self._snapshot(), before) + + def test_check_fails_read_only_when_outputs_are_only_missing(self) -> None: + self._write_shader("op", "op") + before = self._snapshot() + + result, output = self._run("--check") + + self.assertEqual(result, 1) + self.assertIn("Missing embedded WGSL headers", output) + self.assertEqual(self._snapshot(), before) + + def test_check_catches_template_syntax_error_without_writing(self) -> None: + self._write_template( + "a", "syntax", "$if :\n " + self._VALID_SHADER, ["syntax"] + ) + before = self._snapshot() + + result, output = self._run("--check") + + self.assertEqual(result, 1) + self.assertIn("runtime/ops/a/syntax.wgsl", output) + self.assertEqual(self._snapshot(), before) + + def test_check_catches_template_name_error_without_writing(self) -> None: + self._write_template( + "op", "name", "$if MISSING:\n " + self._VALID_SHADER, ["name"] + ) + before = self._snapshot() + + result, output = self._run("--check") + + self.assertEqual(result, 1) + self.assertIn("runtime/ops/op/name.wgsl", output) + self.assertEqual(self._snapshot(), before) + + def test_interrupted_commit_is_detected_and_repaired(self) -> None: + self._write_shader("op", "op") + before = self._snapshot() + real_replace = os.replace + + def interrupt_second(source, destination): + interrupt_second.calls += 1 + if interrupt_second.calls == 2: + raise KeyboardInterrupt("injected interruption") + return real_replace(source, destination) + + interrupt_second.calls = 0 + with mock.patch("os.replace", side_effect=interrupt_second): + with self.assertRaises(KeyboardInterrupt): + self._run() + + self.assertEqual(self._snapshot(), before) + self._assert_no_temps() + check_result, check_output = self._run("--check") + self.assertEqual(check_result, 1) + self.assertNotIn("Orphan", check_output) + + normal_result, _ = self._run() + self.assertEqual(normal_result, 0) + final_check_result, _ = self._run("--check") + self.assertEqual(final_check_result, 0) + self._assert_no_temps() + + def test_interrupt_after_replace_restores_tree(self) -> None: + self._write_shader("op", "op") + before = self._snapshot() + real_replace = os.replace + + def interrupt_after_second(source, destination): + interrupt_after_second.calls += 1 + result = real_replace(source, destination) + if interrupt_after_second.calls == 2: + raise KeyboardInterrupt("injected post-replace interruption") + return result + + interrupt_after_second.calls = 0 + with mock.patch("os.replace", side_effect=interrupt_after_second): + with self.assertRaises(KeyboardInterrupt): + self._run() + + self.assertEqual(self._snapshot(), before) + self._assert_no_temps() + + def test_generation_renders_once_and_second_run_does_no_io(self) -> None: + shaders = [ + self._write_shader("a", "first"), + self._write_shader("b", "second"), + ] + render_counts = {shader: 0 for shader in shaders} + real_headers_for_shader = g.headers_for_shader + + def counted(shader): + render_counts[shader] += 1 + return real_headers_for_shader(shader) + + with mock.patch.object(g, "headers_for_shader", side_effect=counted): + first_result, _ = self._run() + self.assertEqual(first_result, 0) + self.assertEqual(render_counts, {shader: 1 for shader in shaders}) + + with mock.patch( + "tempfile.mkstemp", wraps=tempfile.mkstemp + ) as mkstemp, mock.patch("os.replace", wraps=os.replace) as replace: + second_result, _ = self._run() + self.assertEqual(second_result, 0) + mkstemp.assert_not_called() + replace.assert_not_called() + + class WgslTemplateEngineTest(unittest.TestCase): """Coverage for the $-block template engine + DTYPE/VEC variant matrix.""" @@ -444,6 +921,234 @@ def test_rms_norm_template_roundtrip_byte_identical(self) -> None: got, want, f"{header_name} not reproduced from rms_norm.wgsl template" ) + def test_to_copy_convert_template_roundtrip_byte_identical(self) -> None: + to_copy_dir = g.BACKEND_ROOT / "runtime/ops/to_copy" + template_path = to_copy_dir / "to_copy_convert.wgsl" + spec = g.parse_template_spec(template_path.with_suffix(".yaml")) + variants = {params["NAME"]: params for params in spec[template_path.stem]} + expected = { + "to_copy_float_to_int": ( + "f32", + "i32", + "c331e00e3171eecbe6317ac9df0a5f9cd6d25da26a9a587250f1cc6086dc3c8f", + ), + "to_copy_int_to_float": ( + "i32", + "f32", + "e18dd733a3838f83eded4977a2a2b21119099c8409b234f12474fae5acc9b195", + ), + } + self.assertEqual(set(variants), set(expected)) + template = template_path.read_text() + + for name, (in_type, out_type, expected_hash) in expected.items(): + params = variants[name] + self.assertEqual( + (params["IN_TYPE"], params["OUT_TYPE"]), (in_type, out_type) + ) + expanded = g.preprocess(template, {**g.WGSL_HELPERS, **params}) + self.assertEqual(g.wgsl_sha256(expanded), expected_hash) + + header = (to_copy_dir / f"{name}_wgsl.h").read_text() + body = header.split('R"(', 1)[1].split(')";', 1)[0][1:] + self.assertEqual(body, expanded) + self.assertEqual(g.embedded_sha256(header), expected_hash) + self.assertEqual(g.parse_workgroup_size(body), (64, 1, 1)) + + entries = {entry.name: entry for entry in g.registry_entries()} + self.assertEqual( + entries["to_copy_float_to_int"].include, + "runtime/ops/to_copy/to_copy_float_to_int_wgsl.h", + ) + self.assertEqual( + entries["to_copy_int_to_float"].include, + "runtime/ops/to_copy/to_copy_int_to_float_wgsl.h", + ) + + def test_extrema_template_roundtrip_byte_identical(self) -> None: + extrema_dir = g.BACKEND_ROOT / "runtime/ops/extrema" + template_path = extrema_dir / "extrema.wgsl" + spec = g.parse_template_spec(template_path.with_suffix(".yaml")) + variants = {params["NAME"]: params for params in spec[template_path.stem]} + expected = { + "amax": ( + "max", + "35fc059d7c72caa17f9cb1128823ecfd8f75be4ce24b6cd4f9629a97b52f64c0", + ), + "amin": ( + "min", + "8cb6035ae4d34eb2a6cc973d93d9847905722e967239c96033fccfe3a1943cb2", + ), + } + self.assertEqual(set(variants), set(expected)) + template = template_path.read_text() + + for name, (reduce_fn, expected_hash) in expected.items(): + params = variants[name] + self.assertEqual(params["REDUCE_FN"], reduce_fn) + expanded = g.preprocess(template, {**g.WGSL_HELPERS, **params}) + self.assertEqual(g.wgsl_sha256(expanded), expected_hash) + + header_path = extrema_dir / f"{name}_wgsl.h" + header = header_path.read_text() + body = header.split('R"(', 1)[1].split(')";', 1)[0][1:] + self.assertEqual(body, expanded) + self.assertEqual(g.embedded_sha256(header), expected_hash) + self.assertEqual(g.parse_workgroup_size(body), (256, 1, 1)) + + entries = {entry.name: entry for entry in g.registry_entries()} + for name in expected: + self.assertEqual( + entries[name].include, + f"runtime/ops/extrema/{name}_wgsl.h", + ) + self.assertEqual(entries[name].symbol, g.symbol_base(name)) + + handler_hashes = { + "amax": "57f929b9f3087dc32403c3587884ce2ed4be2d03c4e80ff7035428b52e7e0e51", + "amin": "5dc947d4781a67df953b9c5970c4ab2119317b29657f983a9d308dfdf123dede", + } + for name, expected_hash in handler_hashes.items(): + handler = g.BACKEND_ROOT / f"runtime/ops/{name}/Reduce.cpp" + self.assertEqual( + hashlib.sha256(handler.read_bytes()).hexdigest(), expected_hash + ) + + def test_logical_binary_template_roundtrip_byte_identical(self) -> None: + logical_dir = g.BACKEND_ROOT / "runtime/ops/logical_binary" + template_path = logical_dir / "logical_binary.wgsl" + spec = g.parse_template_spec(template_path.with_suffix(".yaml")) + variants = {params["NAME"]: params for params in spec[template_path.stem]} + expected = { + "logical_and": ( + "&", + "cf7c1d1dbba94e429120796c9c25a6717786cca03c08f3bd1e291d5627089c20", + ), + "logical_or": ( + "|", + "4ad19ee04e2c7b396b4669cf44f95133d658c3ec2e6f37d7b271bedc0e582ecf", + ), + } + self.assertEqual(set(variants), set(expected)) + template = template_path.read_text() + + for name, (op, expected_hash) in expected.items(): + params = variants[name] + self.assertEqual(params["OP"], op) + expanded = g.preprocess(template, {**g.WGSL_HELPERS, **params}) + self.assertEqual(g.wgsl_sha256(expanded), expected_hash) + + header_path = logical_dir / f"{name}_wgsl.h" + header = header_path.read_text() + body = header.split('R"(', 1)[1].split(')";', 1)[0][1:] + self.assertEqual(body, expanded) + self.assertEqual(g.embedded_sha256(header), expected_hash) + self.assertEqual(g.parse_workgroup_size(body), (64, 1, 1)) + + entries = {entry.name: entry for entry in g.registry_entries()} + for name in expected: + self.assertEqual( + entries[name].include, + f"runtime/ops/logical_binary/{name}_wgsl.h", + ) + self.assertEqual(entries[name].symbol, g.symbol_base(name)) + + handler_hashes = { + "logical_and": "eb85a8f97ee7640298a661da49feb08aa79b8c24d3d4458b71d24d3f01bc388d", + "logical_or": "bda18617f7077fee5a812c21cdc495c89542a1688f7e1ef6739ed01da343a66b", + } + for name, expected_hash in handler_hashes.items(): + handler = ( + g.BACKEND_ROOT / f"runtime/ops/{name}/Logical{name[8:].title()}.cpp" + ) + self.assertEqual( + hashlib.sha256(handler.read_bytes()).hexdigest(), expected_hash + ) + + def test_binary_family_roundtrip_byte_identical(self) -> None: + binary_dir = g.BACKEND_ROOT / "runtime/ops/binary_op" + template_path = binary_dir / "binary_op.wgsl" + spec = g.parse_template_spec(template_path.with_suffix(".yaml")) + variants = {params["NAME"]: params for params in spec[template_path.stem]} + expected = { + "binary_div": ( + 0, + "e36b560fd623dd5337b9ae57acd8981c9c635b995d6021caf1331c182cd3f0cd", + ), + "binary_sub": ( + 0, + "63209ff70422a21fc340d9aadba0945bc259bba89bdf05db018a6507d01c7ae5", + ), + "binary_minimum": ( + 1, + "929b7ba85936e3652baea9f4e5e7f049d232c7ae7a74814a536b4c2674897972", + ), + "binary_pow": ( + 1, + "a88c161bd3f43d21a72ebd8ca6f8611b6b9b854e3572a8e6b820602091bc464c", + ), + "binary_floor_divide": ( + 1, + "baf71d277da79389315a6b96b439e7f0a55842e8288283f2af121f84536b3af3", + ), + "binary_mul": ( + 1, + "d248c0f1856b57115a5001a47f4936caa564dd3b787c02ceba504a13ab987812", + ), + } + self.assertEqual(set(variants), set(expected)) + template = template_path.read_text() + entries = {entry.name: entry for entry in g.registry_entries()} + + for name, (inline, expected_hash) in expected.items(): + params = variants[name] + self.assertEqual(params["INLINE"], inline) + expanded = g.preprocess(template, {**g.WGSL_HELPERS, **params}) + self.assertEqual(g.wgsl_sha256(expanded), expected_hash) + + header = (binary_dir / f"{name}_wgsl.h").read_text() + literal = header.split('R"(', 1)[1].split(')";', 1)[0] + self.assertEqual(literal, "\n" + expanded) + self.assertEqual(g.embedded_sha256(header), expected_hash) + self.assertEqual(g.parse_workgroup_size(expanded), (64, 1, 1)) + self.assertIn(f"k{g.symbol_base(name)}WGSL", header) + self.assertEqual( + entries[name].include, + f"runtime/ops/binary_op/{name}_wgsl.h", + ) + + def test_unary_template_roundtrip_byte_identical(self) -> None: + unary_dir = g.BACKEND_ROOT / "runtime/ops/unary" + template_path = unary_dir / "unary.wgsl" + spec = g.parse_template_spec(template_path.with_suffix(".yaml")) + variants = {params["NAME"]: params for params in spec[template_path.stem]} + expected = { + "abs": "39d3c163fdf6a92286828f4b3217e00294e3ca5634a878ed5fd34e3b1cdf0a27", + "cos": "9df78873e5fae98d347c26db2a02b047ea3d5d2c93f0761cb9ac6995f9a71ab2", + "exp": "3171399bc36acf9c1cb2a03c2a31038318203c4c63ab03c4881df7a660346020", + "hardswish": "c874a15ef6cdaec71187296016cc2a1515f5e7c889b97dfa8fd4b278e6e2c3d5", + "neg": "8851b9f42d14153f6f04484fee2f8bf67bda26dea892ff48768e09e6ad49cee1", + "round": "8f3e0edbeb81aa50f35e691c78554e8057fa8d78fe8a86454f4f42e5e8871452", + "rsqrt": "108765d5a23b87473f34651875d08abf2a5fa8980bd92fc8cbe3617295097747", + "sin": "e5762804773659d348fddddcef4935807ae6fe7d92c92eb17a2f44aae8f2c5b9", + "sqrt": "008534ae365969f5c180b42e8d6d0b131df78f181e5435abbcafc3ffb8be8aac", + "tanh": "5bd7eb1c6411940d84a9b311884f35b39f15b82103b14bab02902290ed6b0339", + } + self.assertEqual(set(variants), set(expected)) + template = template_path.read_text() + entries = {entry.name: entry for entry in g.registry_entries()} + for name, expected_hash in expected.items(): + expanded = g.preprocess(template, {**g.WGSL_HELPERS, **variants[name]}) + self.assertEqual(g.wgsl_sha256(expanded), expected_hash) + header = (unary_dir / f"{name}_wgsl.h").read_text() + body = header.split('R"(', 1)[1].split(')";', 1)[0][1:] + self.assertEqual(body, expanded) + self.assertEqual(g.embedded_sha256(header), expected_hash) + self.assertEqual(g.parse_workgroup_size(body), (256, 1, 1)) + self.assertEqual(entries[name].include, f"runtime/ops/unary/{name}_wgsl.h") + self.assertEqual(entries[name].symbol, g.symbol_base(name)) + self.assertTrue({"clamp", "pow_scalar"}.isdisjoint(variants)) + def test_rms_norm_half_variant_is_type_correct(self) -> None: # A DTYPE=half expansion must emit compilable WGSL: `enable f16;`, an f32 # accumulator, loads widened to f32 for the reduction, and the store From a17d30415ca1d2bcec6ebb063619c22a4e04b507 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 5 Aug 2026 12:54:54 -0700 Subject: [PATCH 2/4] Update [ghstack-poisoned] --- .../webgpu/runtime/WebGPUShaderRegistry.cpp | 10 +- .../webgpu/runtime/ops/conv1d_dw/Conv1dDW.cpp | 237 ++++++++++++++++-- .../webgpu/runtime/ops/conv1d_dw/conv1d.wgsl | 53 ++++ .../runtime/ops/conv1d_dw/conv1d_wgsl.h | 77 ++++++ backends/webgpu/test/test_wgsl_codegen.py | 6 +- 5 files changed, 355 insertions(+), 28 deletions(-) create mode 100644 backends/webgpu/runtime/ops/conv1d_dw/conv1d.wgsl create mode 100644 backends/webgpu/runtime/ops/conv1d_dw/conv1d_wgsl.h diff --git a/backends/webgpu/runtime/WebGPUShaderRegistry.cpp b/backends/webgpu/runtime/WebGPUShaderRegistry.cpp index 25de35d49d5..480944ea93d 100644 --- a/backends/webgpu/runtime/WebGPUShaderRegistry.cpp +++ b/backends/webgpu/runtime/WebGPUShaderRegistry.cpp @@ -38,6 +38,7 @@ #include #include #include +#include #include #include #include @@ -152,7 +153,7 @@ namespace executorch::backends::webgpu { namespace { -constexpr std::array kShaderRegistry = {{ +constexpr std::array kShaderRegistry = {{ { "abs", kAbsWGSL, @@ -363,6 +364,13 @@ constexpr std::array kShaderRegistry = {{ kConstantPadNdWorkgroupSizeY, kConstantPadNdWorkgroupSizeZ, }, + { + "conv1d", + kConv1dWGSL, + kConv1dWorkgroupSizeX, + kConv1dWorkgroupSizeY, + kConv1dWorkgroupSizeZ, + }, { "conv1d_dw", kConv1dDwWGSL, diff --git a/backends/webgpu/runtime/ops/conv1d_dw/Conv1dDW.cpp b/backends/webgpu/runtime/ops/conv1d_dw/Conv1dDW.cpp index ae86ccffba8..64d8c2d2380 100644 --- a/backends/webgpu/runtime/ops/conv1d_dw/Conv1dDW.cpp +++ b/backends/webgpu/runtime/ops/conv1d_dw/Conv1dDW.cpp @@ -12,6 +12,7 @@ #include #include #include +#include #include @@ -49,8 +50,24 @@ uint32_t conv1d_out_len( int64_t stride, int64_t padding, int64_t dilation) { - return static_cast( - (in_len + 2 * padding - dilation * (k - 1) - 1) / stride + 1); + if (in_len <= 0 || k <= 0 || stride <= 0 || padding < 0 || dilation <= 0) { + throw std::runtime_error("conv1d: invalid geometry parameter"); + } + constexpr int64_t kMaxShaderIndex = std::numeric_limits::max(); + if (in_len > kMaxShaderIndex || k > kMaxShaderIndex || + stride > kMaxShaderIndex || padding > kMaxShaderIndex || + dilation > kMaxShaderIndex) { + throw std::runtime_error("conv1d: geometry parameter exceeds i32"); + } + const int64_t numerator = in_len + 2 * padding - dilation * (k - 1) - 1; + if (numerator < 0) { + throw std::runtime_error("conv1d: kernel exceeds padded input"); + } + const int64_t out_len = numerator / stride + 1; + if (static_cast(out_len) > UINT32_MAX) { + throw std::runtime_error("conv1d: output length exceeds u32"); + } + return static_cast(out_len); } int64_t first_int(const std::vector& v) { @@ -71,6 +88,24 @@ static_assert( sizeof(Conv1dPwParams) == 32, "Conv1dPwParams must match the WGSL Params struct (32 bytes)"); +struct Conv1dParams { + uint32_t in_channels; + uint32_t out_channels; + uint32_t in_len; + uint32_t out_len; + uint32_t kernel_size; + uint32_t stride; + uint32_t padding; + uint32_t dilation; + uint32_t numel; + uint32_t has_bias; +}; +static_assert( + sizeof(Conv1dParams) == 40, + "Conv1dParams must match the WGSL Params struct (40 bytes)"); +constexpr uint64_t kMaxConv1dDispatchElements = + static_cast(std::numeric_limits::max()); + // Pointwise conv1d (K=1, groups=1): a per-position matmul over channels. void add_conv1d_pw_node( WebGPUGraph& graph, @@ -203,6 +238,154 @@ void add_conv1d_pw_node( graph.own_uniform_buffer(params_buf); } +// General groups=1 conv1d. Voxtral uses K=3 with stride 1 then 2. +void add_conv1d_node( + WebGPUGraph& graph, + int in_id, + int weight_id, + int bias_id, + int out_id, + uint32_t stride, + uint32_t padding, + uint32_t dilation) { + WGPUDevice device = graph.device(); + const auto& in = graph.get_tensor(in_id); + const auto& weight = graph.get_tensor(weight_id); + const auto& out = graph.get_tensor(out_id); + const bool has_bias = + graph.get_value_type(bias_id) == WebGPUGraph::ValueType::Tensor; + if (!utils::is_fp32_tensor(in) || !utils::is_fp32_tensor(weight) || + !utils::is_fp32_tensor(out)) { + throw std::runtime_error("conv1d: input, weight, and output must be fp32"); + } + + const uint32_t expected_out_len = conv1d_out_len( + in.dims.at(2), weight.dims.at(2), stride, padding, dilation); + const uint32_t batch = static_cast(in.dims.at(0)); + const uint32_t in_channels = static_cast(in.dims.at(1)); + const uint32_t in_len = static_cast(in.dims.at(2)); + const uint32_t out_channels = static_cast(out.dims.at(1)); + const uint32_t out_len = static_cast(out.dims.at(2)); + const uint32_t kernel_size = static_cast(weight.dims.at(2)); + if (out.dims.at(0) != in.dims.at(0) || out_len != expected_out_len || + weight.dims.at(0) != out.dims.at(1) || + weight.dims.at(1) != in.dims.at(1)) { + throw std::runtime_error("conv1d: shape mismatch"); + } + + const uint64_t in_numel = utils::check_fp32(in, "conv1d", "input"); + const uint64_t out_numel = utils::check_fp32(out, "conv1d", "output"); + const uint64_t weight_numel = utils::check_fp32(weight, "conv1d", "weight"); + if (in_numel != static_cast(batch) * in_channels * in_len || + out_numel != static_cast(batch) * out_channels * out_len || + weight_numel != + static_cast(out_channels) * in_channels * kernel_size || + in_numel > UINT32_MAX || weight_numel > UINT32_MAX || + out_numel > kMaxConv1dDispatchElements) { + throw std::runtime_error("conv1d: fp32 byte-size or u32 mismatch"); + } + if (has_bias) { + const auto& bias = graph.get_tensor(bias_id); + if (!utils::is_fp32_tensor(bias) || bias.dims.size() != 1 || + bias.dims.at(0) != out.dims.at(1) || + utils::check_fp32(bias, "conv1d", "bias") != out_channels) { + throw std::runtime_error("conv1d: bias shape mismatch"); + } + } + + Conv1dParams params = {}; + params.in_channels = in_channels; + params.out_channels = out_channels; + params.in_len = in_len; + params.out_len = out_len; + params.kernel_size = kernel_size; + params.stride = stride; + params.padding = padding; + params.dilation = dilation; + params.numel = static_cast(out_numel); + params.has_bias = has_bias ? 1u : 0u; + + const uint32_t wg_size = + utils::clamp_workgroup_size(device, kConv1dWorkgroupSizeX); + const utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( + device, params.numel, wg_size, "conv1d"); + WGPUConstantEntry wg_size_constant = utils::make_wg_size_constant(wg_size); + WGPUBuffer params_buf = graph.create_params_buffer(params); + WGPUBuffer bias_buf = + has_bias ? graph.get_tensor(bias_id).buffer : weight.buffer; + const uint64_t bias_size = + has_bias ? graph.get_tensor(bias_id).nbytes : weight.nbytes; + + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kConv1dWGSL, + { + {0, WGPUBufferBindingType_ReadOnlyStorage, in.buffer, in.nbytes}, + {1, WGPUBufferBindingType_Storage, out.buffer, out.nbytes}, + {2, + WGPUBufferBindingType_ReadOnlyStorage, + weight.buffer, + weight.nbytes}, + {3, WGPUBufferBindingType_ReadOnlyStorage, bias_buf, bias_size}, + {4, WGPUBufferBindingType_Uniform, params_buf, sizeof(Conv1dParams)}, + }, + &wg_size_constant, + 1); + const size_t dispatch_idx = graph.add_dispatch( + {bundle.pipeline, + bundle.bind_group, + workgroup_count.x, + "conv1d", + workgroup_count.y}); + + graph.add_tensor_resize_hook( + in_id, + [in_id, + out_id, + in_channels, + out_channels, + kernel_size, + stride, + padding, + dilation, + has_bias, + wg_size, + dispatch_idx, + params_buf](WebGPUGraph& g) { + const auto& dims = g.cur_dims(in_id); + if (dims.size() != 3 || dims[0] <= 0 || dims[1] <= 0 || dims[2] <= 0 || + dims[1] != static_cast(in_channels)) { + throw std::runtime_error("conv1d(resize): input shape changed"); + } + Conv1dParams p = {}; + p.in_channels = in_channels; + p.out_channels = out_channels; + p.in_len = static_cast(dims[2]); + p.out_len = + conv1d_out_len(dims[2], kernel_size, stride, padding, dilation); + p.kernel_size = kernel_size; + p.stride = stride; + p.padding = padding; + p.dilation = dilation; + const uint64_t input_numel = utils::numel(dims); + const uint64_t numel = utils::numel( + {dims[0], out_channels, static_cast(p.out_len)}); + if (input_numel > UINT32_MAX || numel > kMaxConv1dDispatchElements) { + throw std::runtime_error( + "conv1d(resize): tensor numel exceeds shader index range"); + } + p.numel = static_cast(numel); + p.has_bias = has_bias ? 1u : 0u; + const utils::WgCount wgc = utils::compute_2d_workgroup_count( + g.device(), p.numel, wg_size, "conv1d(resize)"); + g.set_cur_dims( + out_id, {dims[0], out_channels, static_cast(p.out_len)}); + wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); + g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x; + g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y; + }); +} + // depthwise-conv1d (groups==C); mirrors Vulkan conv1d_dw (Convolution.cpp:755). void convolution_impl(WebGPUGraph& graph, const std::vector& args) { // args mirror Vulkan conv1d_dw; bias (arg 2) may be Null; out=args.back(). @@ -242,38 +425,44 @@ void convolution_impl(WebGPUGraph& graph, const std::vector& args) { const bool transposed = graph.get_bool(transposed_id); const int64_t groups = graph.get_int(groups_id); - // Pointwise (K=1, groups=1): a matmul over channels; stride-1 / no-pad only. - if (!transposed && groups == 1 && weight_tensor.dims.at(2) == 1 && - first_int(graph.get_int_list(stride_id)) == 1 && - first_int(graph.get_int_list(padding_id)) == 0) { - add_conv1d_pw_node(graph, in_id, weight_id, bias_id, out_id); - return; - } - - // Otherwise only the depthwise config (groups==C, weight [C,1,K]). - if (transposed || groups != static_cast(channels) || - weight_tensor.dims.at(0) != static_cast(channels) || - weight_tensor.dims.at(1) != 1) { - throw std::runtime_error( - "convolution: only depthwise or pointwise conv1d supported"); - } - const int64_t stride_i = first_int(graph.get_int_list(stride_id)); const int64_t padding_i = first_int(graph.get_int_list(padding_id)); const int64_t dilation_i = first_int(graph.get_int_list(dilation_id)); - if (stride_i < 1) { - throw std::runtime_error("convolution: stride must be >= 1"); + if (stride_i < 1 || stride_i > std::numeric_limits::max()) { + throw std::runtime_error("convolution: stride must fit positive i32"); } - if (padding_i < 0) { - throw std::runtime_error("convolution: padding must be >= 0"); + if (padding_i < 0 || padding_i > std::numeric_limits::max()) { + throw std::runtime_error("convolution: padding must fit nonnegative i32"); } - if (dilation_i < 1) { - throw std::runtime_error("convolution: dilation must be >= 1"); + if (dilation_i < 1 || dilation_i > std::numeric_limits::max()) { + throw std::runtime_error("convolution: dilation must fit positive i32"); } const uint32_t stride = static_cast(stride_i); const uint32_t padding = static_cast(padding_i); const uint32_t dilation = static_cast(dilation_i); + // Pointwise (K=1, groups=1): a matmul over channels; stride-1 / no-pad only. + if (!transposed && groups == 1 && weight_tensor.dims.at(2) == 1 && + stride_i == 1 && padding_i == 0) { + add_conv1d_pw_node(graph, in_id, weight_id, bias_id, out_id); + return; + } + + const bool is_depthwise = !transposed && + groups == static_cast(channels) && + weight_tensor.dims.at(0) == static_cast(channels) && + weight_tensor.dims.at(1) == 1; + if (!is_depthwise && !transposed && groups == 1) { + add_conv1d_node( + graph, in_id, weight_id, bias_id, out_id, stride, padding, dilation); + return; + } + + if (!is_depthwise) { + throw std::runtime_error( + "convolution: only depthwise, pointwise, or groups=1 conv1d supported"); + } + uint64_t out_numel = 1; for (int64_t d : out_tensor.dims) { out_numel *= static_cast(d); diff --git a/backends/webgpu/runtime/ops/conv1d_dw/conv1d.wgsl b/backends/webgpu/runtime/ops/conv1d_dw/conv1d.wgsl new file mode 100644 index 00000000000..c51594d2005 --- /dev/null +++ b/backends/webgpu/runtime/ops/conv1d_dw/conv1d.wgsl @@ -0,0 +1,53 @@ +override wg_size: u32 = 64u; + +struct Params { + in_channels: u32, + out_channels: u32, + in_len: u32, + out_len: u32, + kernel_size: u32, + stride: u32, + padding: u32, + dilation: u32, + numel: u32, + has_bias: u32, +}; + +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; +@group(0) @binding(2) var weight: array; +@group(0) @binding(3) var bias: array; +@group(0) @binding(4) var params: Params; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.numel) { + return; + } + + let out_t = idx % params.out_len; + let out_c = (idx / params.out_len) % params.out_channels; + let batch = idx / (params.out_channels * params.out_len); + var sum = 0.0; + + for (var in_c = 0u; in_c < params.in_channels; in_c = in_c + 1u) { + for (var k = 0u; k < params.kernel_size; k = k + 1u) { + let in_t = i32(out_t * params.stride + k * params.dilation) - + i32(params.padding); + if (in_t >= 0 && in_t < i32(params.in_len)) { + let input_idx = + (batch * params.in_channels + in_c) * params.in_len + u32(in_t); + let weight_idx = + (out_c * params.in_channels + in_c) * params.kernel_size + k; + sum = fma(input[input_idx], weight[weight_idx], sum); + } + } + } + if (params.has_bias != 0u) { + sum = sum + bias[out_c]; + } + output[idx] = sum; +} diff --git a/backends/webgpu/runtime/ops/conv1d_dw/conv1d_wgsl.h b/backends/webgpu/runtime/ops/conv1d_dw/conv1d_wgsl.h new file mode 100644 index 00000000000..e29af80459b --- /dev/null +++ b/backends/webgpu/runtime/ops/conv1d_dw/conv1d_wgsl.h @@ -0,0 +1,77 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from conv1d.wgsl - DO NOT EDIT. +// wgsl-sha256: 7bc955b7f43473aab96222e7a1228973b65e38e2e6a19e2b9793e0ed7f3768d9 +inline constexpr const char* kConv1dWGSL = R"( +override wg_size: u32 = 64u; + +struct Params { + in_channels: u32, + out_channels: u32, + in_len: u32, + out_len: u32, + kernel_size: u32, + stride: u32, + padding: u32, + dilation: u32, + numel: u32, + has_bias: u32, +}; + +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; +@group(0) @binding(2) var weight: array; +@group(0) @binding(3) var bias: array; +@group(0) @binding(4) var params: Params; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.numel) { + return; + } + + let out_t = idx % params.out_len; + let out_c = (idx / params.out_len) % params.out_channels; + let batch = idx / (params.out_channels * params.out_len); + var sum = 0.0; + + for (var in_c = 0u; in_c < params.in_channels; in_c = in_c + 1u) { + for (var k = 0u; k < params.kernel_size; k = k + 1u) { + let in_t = i32(out_t * params.stride + k * params.dilation) - + i32(params.padding); + if (in_t >= 0 && in_t < i32(params.in_len)) { + let input_idx = + (batch * params.in_channels + in_c) * params.in_len + u32(in_t); + let weight_idx = + (out_c * params.in_channels + in_c) * params.kernel_size + k; + sum = fma(input[input_idx], weight[weight_idx], sum); + } + } + } + if (params.has_bias != 0u) { + sum = sum + bias[out_c]; + } + output[idx] = sum; +} +)"; + +inline constexpr uint32_t kConv1dWorkgroupSizeX = 64; +inline constexpr uint32_t kConv1dWorkgroupSizeY = 1; +inline constexpr uint32_t kConv1dWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/test/test_wgsl_codegen.py b/backends/webgpu/test/test_wgsl_codegen.py index d293867e7ab..9990297b9f0 100644 --- a/backends/webgpu/test/test_wgsl_codegen.py +++ b/backends/webgpu/test/test_wgsl_codegen.py @@ -220,14 +220,14 @@ def test_generated_output_manifest_digest(self) -> None: digest.update(b"\0") digest.update(output.read_bytes()) digest.update(b"\0") - self.assertEqual(len(outputs), 135) + self.assertEqual(len(outputs), 136) self.assertEqual( digest.hexdigest(), - "61bcf7671ba11a1d9f26e6488d7c62de855728e6cef6ef7983e2a04c4e1c79ca", + "0512f8d258952e446ffaedcb653b6a3a720eccf8a6b5327d95fd454a912214a3", ) self.assertEqual( hashlib.sha256(g.registry_path().read_bytes()).hexdigest(), - "b3e84b54e6fcb7e48477018d91de62977202418f3b45cb6c28a70dfe86944176", + "28aaa7a8d3e916df43e407120e91d487d0d51cbc5ca93c56bd822d25d109890e", ) def test_rope_hf_reconstructs_full_2d_grid_stride(self) -> None: From 12a86d3e0e9635ab5d063a6e3a8ca6823b230719 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 5 Aug 2026 13:43:44 -0700 Subject: [PATCH 3/4] [ExecuTorch][WebGPU] Cover Conv1d routes and dynamic manifests Pull Request resolved: https://github.com/pytorch/executorch/pull/21601 **Cover Conv1d routing, resize, and dynamic manifest replay** Static numeric coverage pins the general groups=1 Conv1d kernel while preserving pointwise and depthwise routing. The manifest harness now separates upper-bound export inputs from live runtime inputs so dynamic execution proves a distinct shape through one loaded graph. Key changes: - Dispatch assertions distinguish general, pointwise, and depthwise Conv1d routes. - Numeric cases cover padding, dilation, batching, optional bias, and dynamic resize. - `Case.export_inputs` and `Case.dynamic_shapes` carry export-time samples and constraints while `Case.inputs` remains the live runtime input. - Generator coverage proves export length 10 and runtime length 7. This is coverage-only; runtime kernels are unchanged. Co-authored-with: Claude Code. ghstack-source-id: 411059163 @exported-using-ghexport Differential Revision: [D114936144](https://our.internmc.facebook.com/intern/diff/D114936144/) --- .../test/native/test_compute_dispatch.cpp | 109 ++++++++++++++++++ .../webgpu/test/native/test_dynamic_shape.cpp | 45 ++++++++ backends/webgpu/test/op_tests/cases.py | 55 ++++++++- .../webgpu/test/op_tests/generate_op_tests.py | 8 +- .../webgpu/test/op_tests/test_generator.py | 21 ++++ backends/webgpu/test/op_tests/test_suite.py | 4 + .../test_dynamic_shape_export.py | 39 +++++++ backends/webgpu/test/ops/test_conv1d_pw.py | 85 +++++++++++++- 8 files changed, 361 insertions(+), 5 deletions(-) diff --git a/backends/webgpu/test/native/test_compute_dispatch.cpp b/backends/webgpu/test/native/test_compute_dispatch.cpp index 6ed7229604b..10ad7af1dca 100644 --- a/backends/webgpu/test/native/test_compute_dispatch.cpp +++ b/backends/webgpu/test/native/test_compute_dispatch.cpp @@ -153,6 +153,92 @@ void expect_dual_q4_topology(WebGPUGraph& graph) { 1); } +struct Conv1dRouteCase { + const char* name; + std::vector input_dims; + std::vector weight_dims; + std::vector output_dims; + int64_t stride; + int64_t padding; + int64_t dilation; + int64_t groups; + const char* expected_kernel; +}; + +void build_conv1d_route_graph( + WebGPUGraph& graph, + const Conv1dRouteCase& test_case) { + namespace vk = vkgraph; + ::flatbuffers::FlatBufferBuilder fbb; + std::vector<::flatbuffers::Offset> values; + auto add_tensor = [&](const std::vector& dims, int mem_obj_id) { + const int id = static_cast(values.size()); + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::VkTensor, + vk::CreateVkTensorDirect( + fbb, + vk::VkDataType::FLOAT32, + &dims, + /*constant_id=*/-1, + mem_obj_id) + .Union())); + return id; + }; + auto add_int = [&](int64_t value) { + const int id = static_cast(values.size()); + values.push_back(vk::CreateVkValue( + fbb, vk::GraphTypes::Int, vk::CreateInt(fbb, value).Union())); + return id; + }; + auto add_int_list = [&](int64_t value) { + const int id = static_cast(values.size()); + const std::vector items = {value}; + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::IntList, + vk::CreateIntListDirect(fbb, &items).Union())); + return id; + }; + + const int input = add_tensor(test_case.input_dims, 0); + const int weight = add_tensor(test_case.weight_dims, 1); + const int bias = static_cast(values.size()); + values.push_back(vk::CreateVkValue(fbb)); + const int stride = add_int_list(test_case.stride); + const int padding = add_int_list(test_case.padding); + const int dilation = add_int_list(test_case.dilation); + const int transposed = static_cast(values.size()); + values.push_back(vk::CreateVkValue( + fbb, vk::GraphTypes::Bool, vk::CreateBool(fbb, false).Union())); + const int output_padding = add_int_list(0); + const int groups = add_int(test_case.groups); + const int output = add_tensor(test_case.output_dims, 2); + const std::vector args = { + input, + weight, + bias, + stride, + padding, + dilation, + transposed, + output_padding, + groups, + output}; + std::vector<::flatbuffers::Offset> chain; + chain.push_back( + vk::CreateOperatorCallDirect(fbb, 0, "aten.convolution.default", &args)); + const std::vector input_ids = { + static_cast(input), static_cast(weight)}; + const std::vector output_ids = {static_cast(output)}; + const auto root = vk::CreateVkGraphDirect( + fbb, "0", &chain, &values, &input_ids, &output_ids); + vk::FinishVkGraphBuffer(fbb, root); + + graph.set_device(g_device); + graph.build(fbb.GetBufferPointer(), nullptr, 0, nullptr); +} + TEST(WebGPUShaderRegistry, FindsKnownShaderAndRejectsUnknownName) { const WebGPUShaderInfo& sigmoid = get_webgpu_shader_info("sigmoid"); EXPECT_EQ(sigmoid.name, "sigmoid"); @@ -191,6 +277,29 @@ TEST(WebGPUQ4RouteSignal, PreservesStaticAndRecordsBothDynamicSignals) { expect_dual_q4_topology(explicit_graph); } +TEST(WebGPUConv1dRoute, SelectsPointwiseGeneralAndSingleChannelDepthwise) { + const Conv1dRouteCase cases[] = { + {"pointwise", {1, 4, 8}, {6, 4, 1}, {1, 6, 8}, 1, 0, 1, 1, "conv1d_pw"}, + {"general", {1, 4, 10}, {6, 4, 3}, {1, 6, 4}, 2, 0, 1, 1, "conv1d"}, + {"single-channel depthwise", + {1, 1, 8}, + {1, 1, 3}, + {1, 1, 8}, + 1, + 1, + 1, + 1, + "conv1d_dw"}, + }; + for (const Conv1dRouteCase& test_case : cases) { + SCOPED_TRACE(test_case.name); + WebGPUGraph graph; + build_conv1d_route_graph(graph, test_case); + ASSERT_EQ(graph.num_dispatches(), 1); + EXPECT_EQ(graph.dispatch_at(0).kernel_name, test_case.expected_kernel); + } +} + TEST(WebGPUComputeDispatch, PipelineKeyCanonicalizesConstants) { WebGPUComputeDispatchDescriptor first; first.shader_name = "sigmoid"; diff --git a/backends/webgpu/test/native/test_dynamic_shape.cpp b/backends/webgpu/test/native/test_dynamic_shape.cpp index 7b3c1d6b0a7..8e172c7472c 100644 --- a/backends/webgpu/test/native/test_dynamic_shape.cpp +++ b/backends/webgpu/test/native/test_dynamic_shape.cpp @@ -18,6 +18,7 @@ // F dyn_rms_chain (rms(rms(x))) at 3 S -> golden (resize CASCADE, DD-4) // G rms+residual H rms*x I dyn_linear J sdpa_dyn K emb_dyn L rope_dyn // M dyn_sigmoid N dyn_select (select_copy(0,-1), dynamic S) +// O ONE dyn_conv1d graph reused across live input lengths // .pte + goldens from test/ops/dynamic_shape/test_dynamic_shape_export.py. // // Artifacts dir: $WEBGPU_DYNAMIC_SHAPE_DIR, else argv[1], else @@ -127,6 +128,42 @@ void check_s(Module& m, const std::string& prefix, int s) { << " golden.size=" << golden.size() << ")"; } +constexpr int kConv1dInChannels = 3; +constexpr int kConv1dOutChannels = 4; +constexpr int kConv1dKernel = 3; +constexpr int kConv1dStride = 2; +constexpr int kConv1dPadding = 1; +constexpr int kConv1dDilation = 2; + +void check_conv1d(Module& module, int length) { + const std::string prefix = g_dir + "/dyn_conv1d.S" + std::to_string(length); + auto input = read_bin(prefix + ".input.bin"); + auto golden = read_bin(prefix + ".golden.bin"); + ASSERT_EQ(input.size(), static_cast(kConv1dInChannels * length)); + ASSERT_FALSE(golden.empty()); + auto tensor = + make_tensor_ptr({1, kConv1dInChannels, length}, std::move(input)); + auto result = module.forward({EValue(tensor)}); + ASSERT_TRUE( + result.ok() && result.get().size() == 1 && result.get()[0].isTensor()) + << "conv1d length=" << length << " forward failed"; + const auto& output = result.get()[0].toTensor(); + const int output_length = (length + 2 * kConv1dPadding - + kConv1dDilation * (kConv1dKernel - 1) - 1) / + kConv1dStride + + 1; + ASSERT_EQ(output.dim(), 3); + ASSERT_EQ(output.size(0), 1); + ASSERT_EQ(output.size(1), kConv1dOutChannels); + ASSERT_EQ(output.size(2), output_length); + const size_t numel = static_cast(kConv1dOutChannels * output_length); + ASSERT_EQ(static_cast(output.numel()), numel); + std::vector got( + output.const_data_ptr(), output.const_data_ptr() + numel); + const float error = max_err(got, golden); + EXPECT_LT(error, 1e-3f) << "conv1d length=" << length << " max_err=" << error; +} + // Dynamic quantized linear: input [M, kLinK] -> output [M, n]. kLinN is the // register-tiled/bicol config; kLinNShmem (N>=2048) routes to the shmem GEMM. constexpr int kLinK = 64; @@ -893,6 +930,14 @@ TEST(DynamicShape, RmsNormReusedGraph) { } } +TEST(DynamicShape, Conv1dReusedGraph) { + Module module(g_dir + "/dyn_conv1d.pte"); + ASSERT_EQ(module.load_forward(), Error::Ok) << "load dyn_conv1d.pte"; + for (int length : {16, 9, 5, 16}) { + check_conv1d(module, length); + } +} + // C2: grow-only reuse — one loaded rms graph run smallest -> largest, so the // FIRST resize grows the dispatch (every other reuse test starts at MAXS and // only shrinks; this catches a hook with a shrink-only short-circuit). diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index 2c78e9192ca..570f4b2a696 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -53,7 +53,11 @@ CompareModule, ) from executorch.backends.webgpu.test.ops.test_conv1d_dw import Conv1dDWModule -from executorch.backends.webgpu.test.ops.test_conv1d_pw import Conv1dPwModule +from executorch.backends.webgpu.test.ops.test_conv1d_pw import ( + Conv1dModule, + Conv1dPwModule, + GENERAL_CONFIGS as _CONV1D_CONFIGS, +) from executorch.backends.webgpu.test.ops.test_conv_with_clamp import ConvWithClampModule from executorch.backends.webgpu.test.ops.test_flip import FlipModule from executorch.backends.webgpu.test.ops.test_floor_divide import FloorDivideModule @@ -654,6 +658,55 @@ def case(name, C, L, kernel, stride, padding, dilation, bias): case("k3s2p1", 4, 8, 3, 2, 1, 1, True), case("dil2", 3, 10, 3, 1, 2, 2, True), case("k5_nobias", 5, 7, 5, 1, 0, 1, False), + case("single_channel_route", 1, 8, 3, 1, 1, 1, True), + ], + atol=1e-3, + rtol=1e-3, + ) + + +@register_op_test("conv1d") +def _conv1d_suite() -> WebGPUTestSuite: + # General groups=1 NCL conv1d; fp64 oracle. The neighboring pointwise and + # depthwise suites remain routing controls for the two retained fast paths. + def case(name, cfg): + n, ic, oc, length, kernel, stride, padding, dilation, bias = cfg + return Case( + name=name, + construct={ + "in_channels": ic, + "out_channels": oc, + "kernel_size": kernel, + "stride": stride, + "padding": padding, + "dilation": dilation, + "bias": bias, + }, + inputs=((n, ic, length),), + ) + + dynamic_cfg = _CONV1D_CONFIGS["voxtral_stride1"] + n, ic, oc, length, kernel, stride, padding, dilation, bias = dynamic_cfg + dynamic_length = torch.export.Dim("manifest_conv1d_length", min=7, max=length) + return WebGPUTestSuite( + module_factory=Conv1dModule, + cases=[ + *[case(name, cfg) for name, cfg in _CONV1D_CONFIGS.items()], + Case( + name="dynamic_length_10_to_7", + construct={ + "in_channels": ic, + "out_channels": oc, + "kernel_size": kernel, + "stride": stride, + "padding": padding, + "dilation": dilation, + "bias": bias, + }, + export_inputs=((n, ic, length),), + inputs=((n, ic, 7),), + dynamic_shapes=({2: dynamic_length},), + ), ], atol=1e-3, rtol=1e-3, diff --git a/backends/webgpu/test/op_tests/generate_op_tests.py b/backends/webgpu/test/op_tests/generate_op_tests.py index 860f0661913..33b7940a9be 100644 --- a/backends/webgpu/test/op_tests/generate_op_tests.py +++ b/backends/webgpu/test/op_tests/generate_op_tests.py @@ -55,13 +55,17 @@ def _materialize(spec) -> torch.Tensor: def export_case(suite: WebGPUTestSuite, case) -> tuple[torch.nn.Module, tuple, object]: - """Build the module + forward inputs and export to an ExecuTorch program.""" + """Build the module and export it, returning the live runtime inputs.""" module = suite.module_factory(**case.construct) # Seed so an unseeded-randn input is reproducible across generations (the golden uses # the SAME tensor, so this only affects which bytes a case sees, never pass/fail). torch.manual_seed(0) inputs = tuple(_materialize(s) for s in case.inputs) - ep = torch.export.export(module, inputs) + export_inputs = inputs + if case.export_inputs is not None: + torch.manual_seed(0) + export_inputs = tuple(_materialize(s) for s in case.export_inputs) + ep = torch.export.export(module, export_inputs, dynamic_shapes=case.dynamic_shapes) prog = to_edge_transform_and_lower( ep, partitioner=[VulkanPartitioner()] ).to_executorch() diff --git a/backends/webgpu/test/op_tests/test_generator.py b/backends/webgpu/test/op_tests/test_generator.py index 65f765812be..ec4125a2818 100644 --- a/backends/webgpu/test/op_tests/test_generator.py +++ b/backends/webgpu/test/op_tests/test_generator.py @@ -61,6 +61,27 @@ def test_generate_case_writes_artifacts(tmp_path): assert entry["golden"]["output_index"] == 0 +def test_export_case_separates_upper_bound_from_runtime_inputs(monkeypatch): + suite = op_test_registry["conv1d"] + case = next(c for c in suite.cases if c.name == "dynamic_length_10_to_7") + export_shapes = [] + exported_dynamic_shapes = [] + real_export = torch.export.export + + def capture_export(module, inputs, **kwargs): + export_shapes.append(tuple(inputs[0].shape)) + exported_dynamic_shapes.append(kwargs.get("dynamic_shapes")) + return real_export(module, inputs, **kwargs) + + monkeypatch.setattr(torch.export, "export", capture_export) + _module, runtime_inputs, prog = g.export_case(suite, case) + + assert export_shapes == [(1, 4, 10)] + assert exported_dynamic_shapes == [case.dynamic_shapes] + assert tuple(runtime_inputs[0].shape) == (1, 4, 7) + assert g._has_vulkan_delegate(prog) + + def test_generate_manifest(tmp_path): g.generate(str(tmp_path), ops=["add"]) manifest = tmp_path / "manifest.json" diff --git a/backends/webgpu/test/op_tests/test_suite.py b/backends/webgpu/test/op_tests/test_suite.py index f2714125c84..805a12e3a22 100644 --- a/backends/webgpu/test/op_tests/test_suite.py +++ b/backends/webgpu/test/op_tests/test_suite.py @@ -58,6 +58,10 @@ class Case: required: bool = True heavy: bool = False golden_fn: Callable | None = None + # Optional upper-bound inputs and shape constraints for a dynamic export. + # `inputs` remain the live tensors written to the runtime manifest. + export_inputs: tuple[Input, ...] | None = None + dynamic_shapes: object | None = None def __post_init__(self) -> None: # Mirror kQ4gswConfigs: every heavy config is required=False (export-gated, never FAILs on absence). diff --git a/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py b/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py index b302b3c120d..ce7dcd42f49 100644 --- a/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py +++ b/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py @@ -19,6 +19,7 @@ import torch from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.backends.webgpu.test.ops.test_conv1d_pw import Conv1dModule from executorch.exir import to_edge_transform_and_lower from executorch.exir.backend.utils import get_delegates, get_non_lowered_nodes @@ -252,9 +253,40 @@ def _write_goldens(model, prefix: str, out_dir: str, s_values) -> None: print(f" golden {prefix} S={s}") +def export_dynamic_conv1d_cases(out_dir: str) -> None: + """Write one dynamic Conv1d program and live-length runtime fixtures.""" + os.makedirs(out_dir, exist_ok=True) + max_length = 16 + lengths = (max_length, 9, 5) + model = Conv1dModule( + in_channels=3, + out_channels=4, + kernel_size=3, + stride=2, + padding=1, + dilation=2, + bias=True, + ).eval() + length_dim = torch.export.Dim("conv1d_length", min=5, max=max_length) + _export( + model, + (_ramp((1, 3, max_length)),), + {"x": {2: length_dim}}, + os.path.join(out_dir, "dyn_conv1d.pte"), + ) + for length in lengths: + x = _ramp((1, 3, length)) + with torch.no_grad(): + golden = model(x) + prefix = os.path.join(out_dir, f"dyn_conv1d.S{length}") + x.detach().numpy().astype(" None: """Write the dynamic + static .pte's and per-S goldens for the native test.""" os.makedirs(out_dir, exist_ok=True) + export_dynamic_conv1d_cases(out_dir) s_dim = torch.export.Dim("s", min=1, max=MAXS) # 1) Single dynamic rms_norm, graph built at S=MAXS (upper bound). @@ -1536,6 +1568,13 @@ def test_export_dynamic_rms(self) -> None: self.assertTrue(os.path.exists(os.path.join(d, "dyn_rms.pte"))) self.assertTrue(os.path.exists(os.path.join(d, "dyn_rms.S1.golden.bin"))) expected = [ + "dyn_conv1d.pte", + "dyn_conv1d.S16.input.bin", + "dyn_conv1d.S16.golden.bin", + "dyn_conv1d.S9.input.bin", + "dyn_conv1d.S9.golden.bin", + "dyn_conv1d.S5.input.bin", + "dyn_conv1d.S5.golden.bin", "dyn_linear_bk64.pte", "dyn_linear_bk64.S512.input.bin", "dyn_linear_bk64.S512.golden.bin", diff --git a/backends/webgpu/test/ops/test_conv1d_pw.py b/backends/webgpu/test/ops/test_conv1d_pw.py index a0de8988f81..989435a6842 100644 --- a/backends/webgpu/test/ops/test_conv1d_pw.py +++ b/backends/webgpu/test/ops/test_conv1d_pw.py @@ -27,6 +27,17 @@ "batch2": (2, 3, 4, 5, True), } +# name -> batch, in_channels, out_channels, L, kernel, stride, padding, +# dilation, bias +GENERAL_CONFIGS = { + "voxtral_stride1": (1, 4, 6, 10, 3, 1, 0, 1, True), + "voxtral_stride2": (1, 6, 5, 10, 3, 2, 0, 1, True), + "no_bias": (1, 3, 2, 9, 3, 1, 0, 1, False), + "padded": (1, 3, 4, 9, 3, 1, 1, 1, True), + "dilated": (1, 2, 3, 11, 3, 1, 2, 2, True), + "batch2": (2, 3, 4, 8, 3, 2, 1, 1, True), +} + class Conv1dPwModule(torch.nn.Module): def __init__(self, in_channels, out_channels, bias) -> None: @@ -42,6 +53,37 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return self.conv(x) +class Conv1dModule(torch.nn.Module): + def __init__( + self, + in_channels, + out_channels, + kernel_size, + stride, + padding, + dilation, + bias, + ) -> None: + super().__init__() + g = torch.Generator().manual_seed(0) + self.conv = torch.nn.Conv1d( + in_channels, + out_channels, + kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + bias=bias, + ) + with torch.no_grad(): + self.conv.weight.normal_(generator=g) + if bias: + self.conv.bias.normal_(generator=g) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.conv(x) + + def _det_input(shape): g = torch.Generator().manual_seed(1) return torch.randn(*shape, generator=g, dtype=torch.float32) @@ -54,6 +96,14 @@ def _lower(cfg): return to_edge_transform_and_lower(ep, partitioner=[VulkanPartitioner()]) +def _lower_general(cfg, dynamic_shapes=None): + n, ic, oc, length, kernel, stride, padding, dilation, bias = cfg + module = Conv1dModule(ic, oc, kernel, stride, padding, dilation, bias).eval() + inputs = (_det_input((n, ic, length)),) + ep = torch.export.export(module, inputs, dynamic_shapes=dynamic_shapes) + return to_edge_transform_and_lower(ep, partitioner=[VulkanPartitioner()]) + + def _delegated(et) -> bool: return any( d.id == "VulkanBackend" @@ -63,9 +113,19 @@ def _delegated(et) -> bool: def _op_delegated(edge, op_substr: str) -> bool: - # op must be absorbed into the delegate, not left as a top-level CPU-fallback node. + # The op must be absorbed into a delegate: absent from the top-level graph AND + # present inside a lowered submodule reached by an executorch_call_delegate node + # (a bare absence check also passes for an empty graph or a renamed op). + from executorch.exir.lowered_backend_module import get_lowered_submodules + gm = edge.exported_program().graph_module - return all(op_substr not in str(getattr(n, "target", "")) for n in gm.graph.nodes) + if any(op_substr in str(getattr(n, "target", "")) for n in gm.graph.nodes): + return False + return any( + op_substr in str(getattr(dn, "target", "")) + for _, lowered, _ in get_lowered_submodules(gm) + for dn in lowered.original_module.graph_module.graph.nodes + ) class Conv1dPwTest(unittest.TestCase): @@ -82,3 +142,24 @@ def test_export_delegates(self) -> None: _op_delegated(edge, "convolution"), f"conv1d not delegated (fell back to CPU) for {name}", ) + + +class Conv1dTest(unittest.TestCase): + def test_export_delegates(self) -> None: + for name, cfg in GENERAL_CONFIGS.items(): + with self.subTest(name=name): + edge = _lower_general(cfg) + self.assertTrue( + _delegated(edge.to_executorch()), + f"Expected a VulkanBackend delegate (conv1d {name})", + ) + self.assertTrue( + _op_delegated(edge, "convolution"), + f"conv1d not delegated (fell back to CPU) for {name}", + ) + + def test_dynamic_length_export_delegates(self) -> None: + length = torch.export.Dim("conv1d_length", min=5, max=16) + edge = _lower_general(GENERAL_CONFIGS["padded"], dynamic_shapes=({2: length},)) + self.assertTrue(_delegated(edge.to_executorch())) + self.assertTrue(_op_delegated(edge, "convolution")) From 8608c4b6354b737e04a7d84e94a434a54ea03cfc Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 5 Aug 2026 13:43:45 -0700 Subject: [PATCH 4/4] [ExecuTorch][WebGPU] Lift FastVLM dispatch limits and harden graph construction Pull Request resolved: https://github.com/pytorch/executorch/pull/21602 **Run production FastVLM shapes with exception-safe WebGPU graph construction** Production FastVLM tensors exceeded WebGPU one-dimensional dispatch limits, while several parameter buffers were not graph-owned before fallible pipeline construction. The runtime now folds flat dispatches across x/y, owns parameters immediately, rejects ambiguous `BOOL`/byte-integer casts, and compiles timestamp validation only in profiling builds. Key changes: - `GELU` and `expand_copy` use overflow-safe folded dispatch and live resize updates. - `Compare`, `expand_copy`, and `to_copy` create graph-owned parameter buffers before pipeline construction, mirroring Vulkan `ComputeGraph` ownership. - `to_copy` rejects `BOOL`-to/from-`INT8` or `UINT8` before same-width flat-copy routing. - Production builds exclude timestamp-query validation and symbols. - Boundary, failure, and real FastVLM encoder coverage exercise the final combined behavior. Kernel math and below-limit one-dimensional dispatch behavior are unchanged. Co-authored-with: Claude Code. ghstack-source-id: 411059172 @exported-using-ghexport Differential Revision: [D114936147](https://our.internmc.facebook.com/intern/diff/D114936147/) --- backends/webgpu/runtime/WebGPUDispatchMath.h | 2 +- backends/webgpu/runtime/WebGPUGraph.cpp | 12 +--- .../webgpu/runtime/ops/compare/Compare.cpp | 6 +- .../runtime/ops/expand_copy/ExpandCopy.cpp | 38 ++++++++--- .../runtime/ops/expand_copy/expand_copy.wgsl | 6 +- .../ops/expand_copy/expand_copy_wgsl.h | 8 ++- backends/webgpu/runtime/ops/gelu/Gelu.cpp | 52 ++++++++++++--- backends/webgpu/runtime/ops/gelu/gelu.wgsl | 14 +++-- backends/webgpu/runtime/ops/gelu/gelu_wgsl.h | 16 +++-- .../webgpu/runtime/ops/to_copy/ToCopy.cpp | 19 +++--- .../test/native/test_compute_dispatch.cpp | 52 +++++++++++++++ .../webgpu/test/native/test_dynamic_shape.cpp | 63 +++++++++++++++++++ .../webgpu/test/native/test_webgpu_utils.cpp | 8 +++ backends/webgpu/test/op_tests/cases.py | 37 ++++++++++- backends/webgpu/test/op_tests/test_suite.py | 3 +- .../test_dynamic_shape_export.py | 50 +++++++++++++++ backends/webgpu/test/ops/test_conv1d_pw.py | 7 +-- 17 files changed, 324 insertions(+), 69 deletions(-) diff --git a/backends/webgpu/runtime/WebGPUDispatchMath.h b/backends/webgpu/runtime/WebGPUDispatchMath.h index 561bebc6b87..60638b499bb 100644 --- a/backends/webgpu/runtime/WebGPUDispatchMath.h +++ b/backends/webgpu/runtime/WebGPUDispatchMath.h @@ -24,7 +24,7 @@ namespace executorch::backends::webgpu::utils { // Ceiling division for non-negative integers (mirrors Vulkan's utils::div_up). template inline T div_up(T a, T b) { - return (a + b - 1) / b; + return a / b + (a % b != 0); } // Product of a tensor's dims; the same accumulation was duplicated per-op. diff --git a/backends/webgpu/runtime/WebGPUGraph.cpp b/backends/webgpu/runtime/WebGPUGraph.cpp index 04331ee6ca7..71579a2cf30 100644 --- a/backends/webgpu/runtime/WebGPUGraph.cpp +++ b/backends/webgpu/runtime/WebGPUGraph.cpp @@ -1518,17 +1518,10 @@ constexpr uint32_t kRouteK16CausalBound = 1u << 11; constexpr uint32_t kRouteBicolSubgroup = 1u << 12; constexpr uint32_t kRouteQwen3Q16K16 = 1u << 13; constexpr uint32_t kRouteQwen3Q32K16 = 1u << 14; -#endif // WGPU_BACKEND_ENABLE_PROFILING - -// Bench gate: compiled out unless WGPU_BACKEND_ENABLE_PROFILING; then the -// WEBGPU_TIMESTAMP_QUERY env var enables per-pass GPU timestamp queries. bool should_timestamp_query() { -#ifdef WGPU_BACKEND_ENABLE_PROFILING return std::getenv("WEBGPU_TIMESTAMP_QUERY") != nullptr; -#else - return false; -#endif } +#endif // WGPU_BACKEND_ENABLE_PROFILING } // namespace #ifdef WGPU_BACKEND_ENABLE_PROFILING @@ -1739,12 +1732,13 @@ size_t WebGPUGraph::execute(const WebGPUExecutionPlan& plan) { return 1; } - // GPU timestamp queries assume one submit; chunked execute is multi-submit. +#ifdef WGPU_BACKEND_ENABLE_PROFILING if (should_timestamp_query()) { throw std::runtime_error( "WebGPU: WEBGPU_TIMESTAMP_QUERY is incompatible with chunked execute " "(multi-submit); disable chunking to use GPU timestamp queries"); } +#endif // WGPU_BACKEND_ENABLE_PROFILING for (size_t chunk_index = 0; chunk_index < plan.dispatch_chunks.size(); chunk_index++) { diff --git a/backends/webgpu/runtime/ops/compare/Compare.cpp b/backends/webgpu/runtime/ops/compare/Compare.cpp index f47f9491191..1870ee85213 100644 --- a/backends/webgpu/runtime/ops/compare/Compare.cpp +++ b/backends/webgpu/runtime/ops/compare/Compare.cpp @@ -85,9 +85,7 @@ void compare_impl( wg_size_constant.key = {"wg_size", WGPU_STRLEN}; wg_size_constant.value = static_cast(wg_size); - WGPUBuffer uniform_buffer = - utils::make_uniform(device, ¶ms, sizeof(CompareParams)); - graph.add_uniform_buffer_bytes(sizeof(CompareParams)); + WGPUBuffer uniform_buffer = graph.create_params_buffer(params); // out (rw storage) + in1/in2 (ro storage) + params (uniform). utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( @@ -142,8 +140,6 @@ void compare_impl( }; graph.add_tensor_resize_hook(in1_id, resize); graph.add_tensor_resize_hook(in2_id, resize); - - graph.own_uniform_buffer(uniform_buffer); } void eq_op(WebGPUGraph& graph, const std::vector& args) { diff --git a/backends/webgpu/runtime/ops/expand_copy/ExpandCopy.cpp b/backends/webgpu/runtime/ops/expand_copy/ExpandCopy.cpp index 5bfc0fa3bf1..9f6a4e68ba6 100644 --- a/backends/webgpu/runtime/ops/expand_copy/ExpandCopy.cpp +++ b/backends/webgpu/runtime/ops/expand_copy/ExpandCopy.cpp @@ -14,6 +14,7 @@ #include +#include #include namespace executorch::backends::webgpu { @@ -34,6 +35,22 @@ void expand_copy_impl(WebGPUGraph& graph, const std::vector& args) { const auto& in_tensor = graph.get_tensor(in_id); const auto& out_tensor = graph.get_tensor(out_id); + if (graph.get_value_type(args.at(1)) != WebGPUGraph::ValueType::IntList) { + throw std::runtime_error( + "WebGPU expand_copy: dynamic target sizes are unsupported"); + } + for (int64_t target_size : graph.get_int_list(args.at(1))) { + if (target_size == -1) { + throw std::runtime_error( + "WebGPU expand_copy: inferred target sizes are unsupported"); + } + } + if (graph.tensor_has_dynamic_dims(in_id) || + graph.tensor_has_dynamic_dims(out_id)) { + throw std::runtime_error( + "WebGPU expand_copy: dynamic shapes are unsupported"); + } + TensorMeta out_meta; TensorMeta in_meta; fill_tensor_meta(out_tensor, &out_meta); @@ -44,21 +61,24 @@ void expand_copy_impl(WebGPUGraph& graph, const std::vector& args) { throw std::runtime_error( "expand_copy: non-fp32 operand (nbytes != numel*4)"); } + if (out_meta.numel > + static_cast(std::numeric_limits::max())) { + throw std::runtime_error( + "WebGPU expand_copy: element count exceeds the flattened 2D dispatch " + "limit"); + } uint32_t wg_size = utils::clamp_workgroup_size(device, kExpandCopyWorkgroupSizeX); - uint32_t workgroup_count = utils::compute_1d_workgroup_count( + utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( device, out_meta.numel, wg_size, "expand_copy"); WGPUConstantEntry wg_size_constant = {}; wg_size_constant.key = {"wg_size", WGPU_STRLEN}; wg_size_constant.value = static_cast(wg_size); - WGPUBuffer out_meta_buf = - utils::make_uniform(device, &out_meta, sizeof(TensorMeta)); - WGPUBuffer in_meta_buf = - utils::make_uniform(device, &in_meta, sizeof(TensorMeta)); - graph.add_uniform_buffer_bytes(2 * sizeof(TensorMeta)); + WGPUBuffer out_meta_buf = graph.create_params_buffer(out_meta); + WGPUBuffer in_meta_buf = graph.create_params_buffer(in_meta); utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( device, @@ -78,10 +98,8 @@ void expand_copy_impl(WebGPUGraph& graph, const std::vector& args) { &wg_size_constant, 1); - graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count}); - - wgpuBufferRelease(out_meta_buf); - wgpuBufferRelease(in_meta_buf); + graph.add_dispatch_2d( + bundle.pipeline, bundle.bind_group, workgroup_count.x, workgroup_count.y); } } // namespace diff --git a/backends/webgpu/runtime/ops/expand_copy/expand_copy.wgsl b/backends/webgpu/runtime/ops/expand_copy/expand_copy.wgsl index 053311a69f4..fab4df15a90 100644 --- a/backends/webgpu/runtime/ops/expand_copy/expand_copy.wgsl +++ b/backends/webgpu/runtime/ops/expand_copy/expand_copy.wgsl @@ -13,8 +13,10 @@ struct TensorMeta { override wg_size: u32 = 64u; @compute @workgroup_size(wg_size, 1, 1) -fn main(@builtin(global_invocation_id) gid: vec3) { - let idx = gid.x; +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); if (idx >= out_meta.numel) { return; } diff --git a/backends/webgpu/runtime/ops/expand_copy/expand_copy_wgsl.h b/backends/webgpu/runtime/ops/expand_copy/expand_copy_wgsl.h index 83c5881604f..f1449f61793 100644 --- a/backends/webgpu/runtime/ops/expand_copy/expand_copy_wgsl.h +++ b/backends/webgpu/runtime/ops/expand_copy/expand_copy_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from expand_copy.wgsl - DO NOT EDIT. -// wgsl-sha256: 99953670bea89e42bc9c689ab80addfd9a331442c8c8f1a5b0c39dbe11c19370 +// wgsl-sha256: b3c032ab961ffde245fc44289b67df3b5e4ca93eedb9ada2f20a3eaa6f10e9c6 inline constexpr const char* kExpandCopyWGSL = R"( @group(0) @binding(0) var input: array; @group(0) @binding(1) var output: array; @@ -30,8 +30,10 @@ struct TensorMeta { override wg_size: u32 = 64u; @compute @workgroup_size(wg_size, 1, 1) -fn main(@builtin(global_invocation_id) gid: vec3) { - let idx = gid.x; +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); if (idx >= out_meta.numel) { return; } diff --git a/backends/webgpu/runtime/ops/gelu/Gelu.cpp b/backends/webgpu/runtime/ops/gelu/Gelu.cpp index 943012515ff..023ed7b5d88 100644 --- a/backends/webgpu/runtime/ops/gelu/Gelu.cpp +++ b/backends/webgpu/runtime/ops/gelu/Gelu.cpp @@ -13,6 +13,7 @@ #include +#include #include #include #include @@ -42,13 +43,18 @@ void gelu_impl(WebGPUGraph& graph, const std::vector& args) { const auto& out_tensor = graph.get_tensor(out_id); utils::check_elementwise_fp32_io(in_tensor, out_tensor, "gelu"); - uint32_t num_elements = - static_cast(out_tensor.nbytes / sizeof(float)); + const uint64_t num_elements64 = out_tensor.nbytes / sizeof(float); + if (num_elements64 > + static_cast(std::numeric_limits::max())) { + throw std::runtime_error( + "WebGPU gelu: element count exceeds the flattened 2D dispatch limit"); + } + const uint32_t num_elements = static_cast(num_elements64); // Each thread handles up to 4 elements (vec4 body + scalar-tail idiom). uint32_t num_vec4_threads = utils::div_up(num_elements, 4u); uint32_t wg_size = utils::clamp_workgroup_size(device, kGeluWorkgroupSizeX); - uint32_t workgroup_count = utils::compute_1d_workgroup_count( + utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( device, num_vec4_threads, wg_size, "gelu"); WGPUConstantEntry wg_constant = utils::make_wg_size_constant(wg_size); @@ -56,9 +62,7 @@ void gelu_impl(WebGPUGraph& graph, const std::vector& args) { GeluParams params = {}; params.num_elements = num_elements; - WGPUBuffer uniform_buffer = - utils::make_uniform(device, ¶ms, sizeof(GeluParams)); - graph.add_uniform_buffer_bytes(sizeof(GeluParams)); + WGPUBuffer uniform_buffer = graph.create_params_buffer(params); // input (read storage) + output (storage) + params. The exact/approximate // choice is baked into the compiled pipeline via the entry point (mirrors @@ -85,10 +89,38 @@ void gelu_impl(WebGPUGraph& graph, const std::vector& args) { 1, exact ? "main_erf" : "main_tanh"); - graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count}); - - // Drop our ref; the bind group keeps the uniform buffer alive until release. - wgpuBufferRelease(uniform_buffer); + const size_t dispatch_idx = graph.add_dispatch_2d( + bundle.pipeline, bundle.bind_group, workgroup_count.x, workgroup_count.y); + + WGPUBuffer params_buf = uniform_buffer; + graph.add_tensor_resize_hook( + in_id, + [in_id, out_id, wg_size, dispatch_idx, params_buf](WebGPUGraph& g) { + const auto& dims = g.cur_dims(in_id); + const uint64_t num_elements64 = utils::numel_of(dims); + if (num_elements64 > + static_cast(std::numeric_limits::max())) { + throw std::runtime_error( + "WebGPU gelu(resize): element count exceeds the flattened 2D " + "dispatch limit"); + } + const uint32_t num_elements = static_cast(num_elements64); + g.set_cur_dims(out_id, dims); + + GeluParams params = {}; + params.num_elements = num_elements; + wgpuQueueWriteBuffer( + g.queue(), params_buf, 0, ¶ms, sizeof(GeluParams)); + + const uint32_t num_vec4_threads = utils::div_up(num_elements, 4u); + const utils::WgCount resized_workgroup_count = + utils::compute_2d_workgroup_count( + g.device(), num_vec4_threads, wg_size, "gelu(resize)"); + g.dispatch_at(dispatch_idx).workgroup_count_x = + resized_workgroup_count.x; + g.dispatch_at(dispatch_idx).workgroup_count_y = + resized_workgroup_count.y; + }); } } // namespace diff --git a/backends/webgpu/runtime/ops/gelu/gelu.wgsl b/backends/webgpu/runtime/ops/gelu/gelu.wgsl index 4f7eb68bc96..9583ef81551 100644 --- a/backends/webgpu/runtime/ops/gelu/gelu.wgsl +++ b/backends/webgpu/runtime/ops/gelu/gelu.wgsl @@ -33,8 +33,11 @@ fn gelu_erf4(x: vec4) -> vec4 { // before use), computes GELU as one vec4 op, then scatters back only the // in-bounds lanes. @compute @workgroup_size(wg_size, 1, 1) -fn main_tanh(@builtin(global_invocation_id) gid: vec3) { - let base = gid.x * 4u; +fn main_tanh( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let thread_idx = gid.x + gid.y * (num_workgroups.x * wg_size); + let base = thread_idx * 4u; if (base >= params.num_elements) { return; } @@ -49,8 +52,11 @@ fn main_tanh(@builtin(global_invocation_id) gid: vec3) { } @compute @workgroup_size(wg_size, 1, 1) -fn main_erf(@builtin(global_invocation_id) gid: vec3) { - let base = gid.x * 4u; +fn main_erf( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let thread_idx = gid.x + gid.y * (num_workgroups.x * wg_size); + let base = thread_idx * 4u; if (base >= params.num_elements) { return; } diff --git a/backends/webgpu/runtime/ops/gelu/gelu_wgsl.h b/backends/webgpu/runtime/ops/gelu/gelu_wgsl.h index 6da12e229af..f8af0f8d2c3 100644 --- a/backends/webgpu/runtime/ops/gelu/gelu_wgsl.h +++ b/backends/webgpu/runtime/ops/gelu/gelu_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from gelu.wgsl - DO NOT EDIT. -// wgsl-sha256: 18f4a82d3bad1ef8703397b871c708804140c4cb382451661f7a77367ac2425f +// wgsl-sha256: 96570753688590fa009ee5503f754cf3eb572dcb3dcae6818220fe06fe3139ee inline constexpr const char* kGeluWGSL = R"( @group(0) @binding(0) var input: array; @group(0) @binding(1) var output: array; @@ -50,8 +50,11 @@ fn gelu_erf4(x: vec4) -> vec4 { // before use), computes GELU as one vec4 op, then scatters back only the // in-bounds lanes. @compute @workgroup_size(wg_size, 1, 1) -fn main_tanh(@builtin(global_invocation_id) gid: vec3) { - let base = gid.x * 4u; +fn main_tanh( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let thread_idx = gid.x + gid.y * (num_workgroups.x * wg_size); + let base = thread_idx * 4u; if (base >= params.num_elements) { return; } @@ -66,8 +69,11 @@ fn main_tanh(@builtin(global_invocation_id) gid: vec3) { } @compute @workgroup_size(wg_size, 1, 1) -fn main_erf(@builtin(global_invocation_id) gid: vec3) { - let base = gid.x * 4u; +fn main_erf( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let thread_idx = gid.x + gid.y * (num_workgroups.x * wg_size); + let base = thread_idx * 4u; if (base >= params.num_elements) { return; } diff --git a/backends/webgpu/runtime/ops/to_copy/ToCopy.cpp b/backends/webgpu/runtime/ops/to_copy/ToCopy.cpp index 3f9c8e29e3c..4fb5fd0d61f 100644 --- a/backends/webgpu/runtime/ops/to_copy/ToCopy.cpp +++ b/backends/webgpu/runtime/ops/to_copy/ToCopy.cpp @@ -70,9 +70,7 @@ void add_convert_op( ConvertParams params = {}; params.num_elements = num_elements; - WGPUBuffer uniform_buffer = - utils::make_uniform(device, ¶ms, sizeof(ConvertParams)); - graph.add_uniform_buffer_bytes(sizeof(ConvertParams)); + WGPUBuffer uniform_buffer = graph.create_params_buffer(params); utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( device, @@ -115,9 +113,6 @@ void add_convert_op( wg_size, "to_copy(resize)"); }); - - // Graph owns it so the resize hook can rewrite it; freed in the dtor. - graph.own_uniform_buffer(uniform_buffer); } // Decode byte-packed bool storage into numeric fp32 values. @@ -160,9 +155,7 @@ void add_bool_to_float_op(WebGPUGraph& graph, int in_id, int out_id) { ConvertParams params = {}; params.num_elements = num_elements; - WGPUBuffer uniform_buffer = - utils::make_uniform(device, ¶ms, sizeof(ConvertParams)); - graph.add_uniform_buffer_bytes(sizeof(ConvertParams)); + WGPUBuffer uniform_buffer = graph.create_params_buffer(params); utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( device, @@ -208,8 +201,6 @@ void add_bool_to_float_op(WebGPUGraph& graph, int in_id, int out_id) { wg_size, "to_copy_bool_to_float(resize)"); }); - - graph.own_uniform_buffer(uniform_buffer); } void to_copy_impl(WebGPUGraph& graph, const std::vector& args) { @@ -223,6 +214,12 @@ void add_to_copy_node(WebGPUGraph& graph, int in_id, int out_id) { const auto& in_tensor = graph.get_tensor(in_id); const auto& out_tensor = graph.get_tensor(out_id); + if (in_tensor.is_bool != out_tensor.is_bool && in_tensor.is_int && + out_tensor.is_int) { + throw std::runtime_error( + "WebGPU to_copy: bool and integer conversions are unsupported"); + } + // Same is_int+width = flat byte copy; unique dtype key in the 32-bit domain. if (in_tensor.is_int == out_tensor.is_int && in_tensor.elem_size == out_tensor.elem_size) { diff --git a/backends/webgpu/test/native/test_compute_dispatch.cpp b/backends/webgpu/test/native/test_compute_dispatch.cpp index 10ad7af1dca..cd6f7514e3b 100644 --- a/backends/webgpu/test/native/test_compute_dispatch.cpp +++ b/backends/webgpu/test/native/test_compute_dispatch.cpp @@ -1008,6 +1008,58 @@ TEST(WebGPURopeValidation, RejectsMalformedGraphsBeforeDispatchAllocation) { } } +TEST(WebGPUToCopyValidation, RejectsBoolAndByteIntegerConversions) { + ASSERT_TRUE(webgpu_operator_registry().has_op("aten._to_copy.default")); + namespace vk = vkgraph; + struct TestCase { + const char* name; + vk::VkDataType input_dtype; + vk::VkDataType output_dtype; + }; + const TestCase cases[] = { + {"bool_to_int8", vk::VkDataType::BOOL, vk::VkDataType::INT8}, + {"bool_to_uint8", vk::VkDataType::BOOL, vk::VkDataType::UINT8}, + {"int8_to_bool", vk::VkDataType::INT8, vk::VkDataType::BOOL}, + {"uint8_to_bool", vk::VkDataType::UINT8, vk::VkDataType::BOOL}, + }; + for (const TestCase& test_case : cases) { + SCOPED_TRACE(test_case.name); + ::flatbuffers::FlatBufferBuilder fbb; + const std::vector dims = {4}; + std::vector<::flatbuffers::Offset> values; + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::VkTensor, + vk::CreateVkTensorDirect(fbb, test_case.input_dtype, &dims, -1, 0) + .Union())); + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::VkTensor, + vk::CreateVkTensorDirect(fbb, test_case.output_dtype, &dims, -1, 1) + .Union())); + const std::vector args = {0, 1}; + std::vector<::flatbuffers::Offset> chain; + chain.push_back( + vk::CreateOperatorCallDirect(fbb, 0, "aten._to_copy.default", &args)); + const std::vector input_ids = {0}; + const std::vector output_ids = {1}; + const auto root = vk::CreateVkGraphDirect( + fbb, "0", &chain, &values, &input_ids, &output_ids); + vk::FinishVkGraphBuffer(fbb, root); + + WebGPUGraph graph; + try { + graph.build(fbb.GetBufferPointer(), nullptr, 0, nullptr); + FAIL() << test_case.name << " unexpectedly built"; + } catch (const std::runtime_error& error) { + EXPECT_STREQ( + error.what(), + "WebGPU to_copy: bool and integer conversions are unsupported"); + } + EXPECT_EQ(graph.memory_stats().num_dispatches, 0); + } +} + TEST(WebGPUExecution, FullySuppressedPlanPerformsNoQueueSubmission) { WebGPUGraph graph; const WebGPUExecutionPlan plan; diff --git a/backends/webgpu/test/native/test_dynamic_shape.cpp b/backends/webgpu/test/native/test_dynamic_shape.cpp index 8e172c7472c..cf16cd7e87c 100644 --- a/backends/webgpu/test/native/test_dynamic_shape.cpp +++ b/backends/webgpu/test/native/test_dynamic_shape.cpp @@ -19,6 +19,7 @@ // G rms+residual H rms*x I dyn_linear J sdpa_dyn K emb_dyn L rope_dyn // M dyn_sigmoid N dyn_select (select_copy(0,-1), dynamic S) // O ONE dyn_conv1d graph reused across live input lengths +// P ONE dyn_gelu graph reused above -> at -> above the old 1D dispatch cap // .pte + goldens from test/ops/dynamic_shape/test_dynamic_shape_export.py. // // Artifacts dir: $WEBGPU_DYNAMIC_SHAPE_DIR, else argv[1], else @@ -34,6 +35,7 @@ #include #include +#include #include #include #include @@ -164,6 +166,38 @@ void check_conv1d(Module& module, int length) { EXPECT_LT(error, 1e-3f) << "conv1d length=" << length << " max_err=" << error; } +constexpr int kGeluOld1dDispatchCap = 4 * 64 * 65535; +constexpr int kGelu2dDispatchBoundary = kGeluOld1dDispatchCap + 1; +constexpr int kGeluPatternSize = 257; + +void check_gelu_2d(Module& module, int elements) { + std::array golden = {}; + std::vector input(static_cast(elements)); + for (int i = 0; i < kGeluPatternSize; i++) { + const float value = -4.0f + 8.0f * i / (kGeluPatternSize - 1); + golden[i] = 0.5f * value * (1.0f + std::erf(value * 0.7071067811865476f)); + } + for (int i = 0; i < elements; i++) { + input[i] = -4.0f + 8.0f * (i % kGeluPatternSize) / (kGeluPatternSize - 1); + } + auto tensor = make_tensor_ptr({elements}, std::move(input)); + auto result = module.forward({EValue(tensor)}); + ASSERT_TRUE( + result.ok() && result.get().size() == 1 && result.get()[0].isTensor()) + << "gelu elements=" << elements << " forward failed"; + const auto& output = result.get()[0].toTensor(); + ASSERT_EQ(output.dim(), 1); + ASSERT_EQ(output.size(0), elements); + ASSERT_EQ(output.numel(), elements); + const float* data = output.const_data_ptr(); + float error = 0.0f; + for (int i = 0; i < elements; i++) { + error = std::fmax(error, std::fabs(data[i] - golden[i % kGeluPatternSize])); + } + EXPECT_LT(error, 1e-4f) << "gelu elements=" << elements + << " max_err=" << error; +} + // Dynamic quantized linear: input [M, kLinK] -> output [M, n]. kLinN is the // register-tiled/bicol config; kLinNShmem (N>=2048) routes to the shmem GEMM. constexpr int kLinK = 64; @@ -938,6 +972,35 @@ TEST(DynamicShape, Conv1dReusedGraph) { } } +TEST(DynamicShape, GeluCrosses2dDispatchBoundary) { + if (std::getenv("WEBGPU_TEST_HEAVY") == nullptr) { + GTEST_SKIP() << "WEBGPU_TEST_HEAVY not set"; + } + Module module(g_dir + "/dyn_gelu_2d.pte"); + ASSERT_EQ(module.load_forward(), Error::Ok) << "load dyn_gelu_2d.pte"; + for (int elements : + {kGelu2dDispatchBoundary, + kGeluOld1dDispatchCap, + kGelu2dDispatchBoundary}) { + check_gelu_2d(module, elements); + } +} + +TEST(DynamicShape, ExpandCopyRejectsDynamicShapesAtLoad) { + const std::string path = g_dir + "/dyn_expand_copy.pte"; + ASSERT_TRUE(std::ifstream(path).good()) << "missing dyn_expand_copy.pte"; + Module module(path); + EXPECT_NE(module.load_forward(), Error::Ok); +} + +TEST(DynamicShape, ExpandCopyRejectsInferredDynamicShapesAtLoad) { + const std::string path = g_dir + "/dyn_expand_copy_inferred.pte"; + ASSERT_TRUE(std::ifstream(path).good()) + << "missing dyn_expand_copy_inferred.pte"; + Module module(path); + EXPECT_NE(module.load_forward(), Error::Ok); +} + // C2: grow-only reuse — one loaded rms graph run smallest -> largest, so the // FIRST resize grows the dispatch (every other reuse test starts at MAXS and // only shrinks; this catches a hook with a shrink-only short-circuit). diff --git a/backends/webgpu/test/native/test_webgpu_utils.cpp b/backends/webgpu/test/native/test_webgpu_utils.cpp index edc0f315294..a839d224b16 100644 --- a/backends/webgpu/test/native/test_webgpu_utils.cpp +++ b/backends/webgpu/test/native/test_webgpu_utils.cpp @@ -14,8 +14,16 @@ #include +#include + using namespace executorch::backends::webgpu; +TEST(WebGPUUtils, DivUpDoesNotOverflowAtUint32Max) { + constexpr uint32_t kMax = std::numeric_limits::max(); + EXPECT_EQ(utils::div_up(kMax, 4u), 1073741824u); + EXPECT_EQ(utils::div_up(kMax, kMax), 1u); +} + TEST(WebGPUUtils, DispatchGridStaysOneDimUnderCeiling) { utils::DispatchGrid g = utils::compute_dispatch_grid_from_limits(1000u, 256u, 65535u, "test"); diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index 570f4b2a696..6a019ab8e9f 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -59,6 +59,10 @@ GENERAL_CONFIGS as _CONV1D_CONFIGS, ) from executorch.backends.webgpu.test.ops.test_conv_with_clamp import ConvWithClampModule +from executorch.backends.webgpu.test.ops.test_expand_copy import ( + CONFIGS as _EXPAND_COPY_CONFIGS, + ExpandCopyModule, +) from executorch.backends.webgpu.test.ops.test_flip import FlipModule from executorch.backends.webgpu.test.ops.test_floor_divide import FloorDivideModule from executorch.backends.webgpu.test.ops.test_grid_priors import GridPriorsModule @@ -667,8 +671,7 @@ def case(name, C, L, kernel, stride, padding, dilation, bias): @register_op_test("conv1d") def _conv1d_suite() -> WebGPUTestSuite: - # General groups=1 NCL conv1d; fp64 oracle. The neighboring pointwise and - # depthwise suites remain routing controls for the two retained fast paths. + # General NCL conv1d; neighboring suites cover the retained fast paths. def case(name, cfg): n, ic, oc, length, kernel, stride, padding, dilation, bias = cfg return Case( @@ -1083,12 +1086,36 @@ def _cat_suite() -> WebGPUTestSuite: N as _GELU_N, ) +_GELU_2D_DISPATCH_BOUNDARY = 4 * 64 * 65535 + 1 +_EXPAND_COPY_2D_DISPATCH_BOUNDARY = 64 * 65535 + 1 + def _gelu_full_range(_shape) -> torch.Tensor: # Reuse the deterministic linspace(-6, 6) spanning negatives/zero/positives. return _gelu_det_input() +@register_op_test("expand_copy") +def _expand_copy_suite() -> WebGPUTestSuite: + cases = [ + Case(name=name, construct={"shape": out_shape}, inputs=(in_shape,)) + for name, (in_shape, out_shape) in _EXPAND_COPY_CONFIGS.items() + ] + cases.append( + Case( + name="dispatch_2d_boundary", + construct={"shape": (_EXPAND_COPY_2D_DISPATCH_BOUNDARY,)}, + inputs=((1,),), + heavy=True, + ) + ) + return WebGPUTestSuite( + module_factory=ExpandCopyModule, + cases=cases, + golden_dtype="float32", + ) + + @register_op_test("gelu") def _gelu_suite() -> WebGPUTestSuite: # erf ("none") is the Florence-2/BART + PyTorch default; tanh is the approx. @@ -1110,6 +1137,12 @@ def _gelu_suite() -> WebGPUTestSuite: construct={"approximate": "none"}, inputs=(InputSpec(shape=(_GELU_N,), gen=_gelu_full_range),), ), + Case( + name="erf_dispatch_2d_boundary", + construct={"approximate": "none"}, + inputs=(InputSpec(shape=(_GELU_2D_DISPATCH_BOUNDARY,), gen="ramp"),), + heavy=True, + ), ], atol=1e-4, rtol=1e-3, diff --git a/backends/webgpu/test/op_tests/test_suite.py b/backends/webgpu/test/op_tests/test_suite.py index 805a12e3a22..17542cd2e55 100644 --- a/backends/webgpu/test/op_tests/test_suite.py +++ b/backends/webgpu/test/op_tests/test_suite.py @@ -58,8 +58,7 @@ class Case: required: bool = True heavy: bool = False golden_fn: Callable | None = None - # Optional upper-bound inputs and shape constraints for a dynamic export. - # `inputs` remain the live tensors written to the runtime manifest. + # Optional upper-bound export inputs; `inputs` stay live manifest tensors. export_inputs: tuple[Input, ...] | None = None dynamic_shapes: object | None = None diff --git a/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py b/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py index ce7dcd42f49..d9f666622ee 100644 --- a/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py +++ b/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py @@ -20,6 +20,7 @@ import torch from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner from executorch.backends.webgpu.test.ops.test_conv1d_pw import Conv1dModule +from executorch.backends.webgpu.test.ops.test_gelu import GeluModule from executorch.exir import to_edge_transform_and_lower from executorch.exir.backend.utils import get_delegates, get_non_lowered_nodes @@ -182,6 +183,20 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return x.select(0, -1) +class DynamicExpandCopyModule(torch.nn.Module): + """Dynamic expand_copy is rejected until its TensorMeta can be resized.""" + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x.expand((4, x.shape[1])).clone() + + +class DynamicExpandCopyInferredModule(torch.nn.Module): + """Dynamic expand_copy whose -1 target hides symbolic provenance.""" + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x.expand((4, -1)).clone() + + def _ramp(shape) -> torch.Tensor: n = 1 for d in shape: @@ -283,10 +298,45 @@ def export_dynamic_conv1d_cases(out_dir: str) -> None: golden.detach().numpy().astype(" None: + """Write a dynamic GELU fixture crossing the old 1D dispatch cap.""" + os.makedirs(out_dir, exist_ok=True) + max_elements = 4 * 64 * 65535 + 1 + model = GeluModule("none").eval() + elements_dim = torch.export.Dim("gelu_elements", min=1024, max=max_elements) + _export( + model, + (torch.empty((max_elements,), dtype=torch.float32),), + {"x": {0: elements_dim}}, + os.path.join(out_dir, "dyn_gelu_2d.pte"), + ) + + +def export_dynamic_expand_copy_rejection_case(out_dir: str) -> None: + """Write a dynamic expand_copy graph that the runtime must reject at load.""" + model = DynamicExpandCopyModule().eval() + elements_dim = torch.export.Dim("expand_elements", min=1, max=8) + _export( + model, + (_ramp((1, 8)),), + {"x": {1: elements_dim}}, + os.path.join(out_dir, "dyn_expand_copy.pte"), + ) + _export( + DynamicExpandCopyInferredModule().eval(), + (_ramp((1, 8)),), + {"x": {1: elements_dim}}, + os.path.join(out_dir, "dyn_expand_copy_inferred.pte"), + ) + + def export_dynamic_shape_cases(out_dir: str) -> None: """Write the dynamic + static .pte's and per-S goldens for the native test.""" os.makedirs(out_dir, exist_ok=True) export_dynamic_conv1d_cases(out_dir) + export_dynamic_expand_copy_rejection_case(out_dir) + if os.environ.get("WEBGPU_TEST_HEAVY"): + export_dynamic_gelu_boundary_cases(out_dir) s_dim = torch.export.Dim("s", min=1, max=MAXS) # 1) Single dynamic rms_norm, graph built at S=MAXS (upper bound). diff --git a/backends/webgpu/test/ops/test_conv1d_pw.py b/backends/webgpu/test/ops/test_conv1d_pw.py index 989435a6842..07036fe4a42 100644 --- a/backends/webgpu/test/ops/test_conv1d_pw.py +++ b/backends/webgpu/test/ops/test_conv1d_pw.py @@ -27,8 +27,7 @@ "batch2": (2, 3, 4, 5, True), } -# name -> batch, in_channels, out_channels, L, kernel, stride, padding, -# dilation, bias +# name -> N, C_in, C_out, L, K, stride, padding, dilation, bias GENERAL_CONFIGS = { "voxtral_stride1": (1, 4, 6, 10, 3, 1, 0, 1, True), "voxtral_stride2": (1, 6, 5, 10, 3, 2, 0, 1, True), @@ -113,9 +112,7 @@ def _delegated(et) -> bool: def _op_delegated(edge, op_substr: str) -> bool: - # The op must be absorbed into a delegate: absent from the top-level graph AND - # present inside a lowered submodule reached by an executorch_call_delegate node - # (a bare absence check also passes for an empty graph or a renamed op). + # Require the op in a delegate, not merely absent from the host graph. from executorch.exir.lowered_backend_module import get_lowered_submodules gm = edge.exported_program().graph_module