From 26ffec7bcbe0eb7cd30ffa28375ee9d44f8ae7c4 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 29 Jul 2026 20:35:01 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- .../ops/embedding_q4gsw/EmbeddingQ4gsw.cpp | 87 +++++++-------- backends/webgpu/runtime/ops/linear/Linear.cpp | 18 ++- .../webgpu/test/native/test_dynamic_shape.cpp | 52 +++++++-- .../test_dynamic_shape_export.py | 105 +++++++++++++++++- 4 files changed, 195 insertions(+), 67 deletions(-) diff --git a/backends/webgpu/runtime/ops/embedding_q4gsw/EmbeddingQ4gsw.cpp b/backends/webgpu/runtime/ops/embedding_q4gsw/EmbeddingQ4gsw.cpp index 9ed4946eab6..b5ea0cdca4c 100644 --- a/backends/webgpu/runtime/ops/embedding_q4gsw/EmbeddingQ4gsw.cpp +++ b/backends/webgpu/runtime/ops/embedding_q4gsw/EmbeddingQ4gsw.cpp @@ -37,17 +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, - uint32_t embed_dim, - uint32_t blocks_per_row, - uint32_t gs_u, - uint32_t groups_per_row, - uint32_t bytes_per_row, + const EmbeddingLayout& layout, uint32_t wg_size, size_t dispatch_idx, WGPUBuffer params_buf) { @@ -56,22 +76,16 @@ void resize_embedding_q4gsw( if (ni == 0) { throw std::runtime_error("WebGPU embedding_q4gsw: zero indices"); } - const uint64_t total_blocks = ni * 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(embed_dim)); + od.push_back(static_cast(layout.embed_dim)); g.set_cur_dims(out_id, od); - EmbeddingParams p = {}; - p.embed_dim = embed_dim; - p.blocks_per_row = blocks_per_row; - p.num_indices = static_cast(ni); - p.group_size = gs_u; - p.groups_per_row = groups_per_row; - p.bytes_per_row = bytes_per_row; - p.total_blocks = static_cast(total_blocks); + 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( @@ -176,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); @@ -230,32 +244,13 @@ void embedding_q4gsw_impl(WebGPUGraph& graph, const std::vector& args) { {bundle.pipeline, bundle.bind_group, workgroup_count, "embedding_q4gsw"}); // Dynamic shapes: recompute counts/dispatch; out = indices + [embed_dim]. - const uint32_t gs_u = static_cast(group_size); WGPUBuffer params_buf = uniform_buffer; graph.add_tensor_resize_hook( indices_id, - [indices_id, - out_id, - embed_dim, - blocks_per_row, - gs_u, - groups_per_row, - bytes_per_row, - wg_size, - dispatch_idx, - params_buf](WebGPUGraph& g) { + [indices_id, out_id, layout, wg_size, dispatch_idx, params_buf]( + WebGPUGraph& g) { resize_embedding_q4gsw( - g, - indices_id, - out_id, - embed_dim, - blocks_per_row, - gs_u, - groups_per_row, - bytes_per_row, - 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 0bb46be111c..4bacc1ff1c8 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; @@ -738,31 +747,31 @@ void run_combined_routes(Module& m, int s) { constexpr int kEmbDim = 64; // Run emb_dyn at N tokens on an already-loaded module (so it can be reused // across N), and compare to the golden. -void run_embedding(Module& m, int n) { - const std::string b = g_dir + "/emb_dyn.S" + std::to_string(n) + "."; +void run_embedding(Module& m, int n, const char* prefix = "emb_dyn") { + const std::string b = g_dir + "/" + prefix + ".S" + std::to_string(n) + "."; std::ifstream f(b + "idx.bin", std::ios::binary | std::ios::ate); - ASSERT_TRUE(f.good()) << "missing emb_dyn.S" << n; + ASSERT_TRUE(f.good()) << "missing " << prefix << ".S" << n; const std::streamsize nb = f.tellg(); - ASSERT_GE(nb, 0) << "missing emb_dyn.S" << n; + ASSERT_GE(nb, 0) << "missing " << prefix << ".S" << n; f.seekg(0); std::vector idx(static_cast(nb) / sizeof(int64_t)); f.read(reinterpret_cast(idx.data()), nb); ASSERT_EQ(idx.size(), static_cast(n)) - << "wrong emb_dyn idx size S" << n; + << "wrong " << prefix << " idx size S" << n; auto golden = read_bin(b + "golden.bin"); auto t = make_tensor_ptr({n}, std::move(idx)); // int64 (Long) host input auto r = m.forward({EValue(t)}); ASSERT_TRUE(r.ok() && !r.get().empty() && r.get()[0].isTensor()) - << "emb N=" << n + << prefix << " N=" << n << " forward failed (err=" << (r.ok() ? 0 : (int)r.error()) << ")"; const auto& out = r.get()[0].toTensor(); const size_t numel = static_cast(n) * kEmbDim; ASSERT_EQ(static_cast(out.numel()), numel) - << "emb N=" << n << " output numel mismatch"; + << prefix << " N=" << n << " output numel mismatch"; std::vector got( out.const_data_ptr(), out.const_data_ptr() + numel); const float e = max_err(got, golden); - EXPECT_LT(e, 5e-3f) << "emb_dyn N=" << n << " max_err=" << e; + EXPECT_LT(e, 5e-3f) << prefix << " N=" << n << " max_err=" << e; } void check_embedding(int n) { @@ -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,6 +1690,16 @@ TEST(DynamicShape, EmbeddingReusedGraph) { } } +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); + } + } +} + // L: dynamic RoPE (two outputs) at several seq-len S. TEST(DynamicShape, Rope) { for (int s : {16, 8, 1}) { 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 1015fe0b8ce..aa3a4cb0d89 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,7 +1353,42 @@ def _export_static_sdpa(out_dir: str, s: int, prefix: str) -> None: EMB_MAXN = 16 -def _export_dynamic_embedding(out_dir: str) -> None: +class _PackedEmbedding(torch.nn.Module): + def __init__(self, is_linear_weight: bool) -> 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, + self.is_linear_weight, + ) + + +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 = model(idx) + idx.detach().numpy().astype(" None: + os.makedirs(out_dir, exist_ok=True) from executorch.backends.webgpu.test.ops.test_embedding_q4gsw import ( _make_quantized_model, _quant_params, @@ -1336,6 +1418,23 @@ def _export_dynamic_embedding(out_dir: str) -> None: ) print(f" golden emb_dyn N={n} (shape {tuple(g.shape)})") + 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. ROPE_NH = 8