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/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/test/native/test_dynamic_shape.cpp b/backends/webgpu/test/native/test_dynamic_shape.cpp index 8c405d441be..2a5a5175680 100644 --- a/backends/webgpu/test/native/test_dynamic_shape.cpp +++ b/backends/webgpu/test/native/test_dynamic_shape.cpp @@ -132,6 +132,7 @@ 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( @@ -232,6 +233,14 @@ void check_linear_tiled(int m_rows) { 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; @@ -935,6 +944,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}) { @@ -1664,12 +1690,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); + } } } 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 a99c9baf46e..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 @@ -297,7 +297,10 @@ 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( @@ -327,7 +330,7 @@ def export_dynamic_shape_cases(out_dir: str) -> None: _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) @@ -370,6 +373,10 @@ def export_dynamic_shape_cases(out_dir: str) -> None: 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 @@ -533,6 +540,46 @@ def _export_dynamic_linear( print(f" golden {prefix} M={m}") +def _export_dynamic_fp32_linear_case( + out_dir: str, + *, + k: int, + bias: bool, + prefix: str, +) -> 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, @@ -1306,52 +1353,64 @@ def _export_static_sdpa(out_dir: str, s: int, prefix: str) -> None: 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, @@ -1370,21 +1429,22 @@ def _export_dynamic_embedding(out_dir: str) -> None: 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.