diff --git a/backends/webgpu/runtime/WebGPUGraph.cpp b/backends/webgpu/runtime/WebGPUGraph.cpp index 5916434ab38..3dafacefe9e 100644 --- a/backends/webgpu/runtime/WebGPUGraph.cpp +++ b/backends/webgpu/runtime/WebGPUGraph.cpp @@ -9,7 +9,6 @@ #include #include #include -#include #include #include @@ -24,6 +23,7 @@ #include #include #include +#include #include #include @@ -165,6 +165,113 @@ std::vector canonical_constants( constexpr const char* kPrepackOpName = "et_vk.prepack.default"; constexpr const char* kQ4gswLinearOpName = "et_vk.linear_q4gsw.default"; constexpr size_t kQ4gswOutputArg = 5; +constexpr const char* kSigmoidOpName = "aten.sigmoid.default"; +constexpr const char* kMulOpName = "aten.mul.Tensor"; + +struct SiluMulParams { + uint32_t num_elements; + uint32_t _pad[3]; +}; + +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; +}; + +bool is_fp32_tensor(const WebGPUTensor& tensor) { + if (tensor.is_int || tensor.elem_size != sizeof(float) || + tensor.buffer == nullptr) { + return false; + } + const uint64_t numel = utils::numel_of(tensor.dims); + return numel <= std::numeric_limits::max() / sizeof(float) && + tensor.nbytes == static_cast(numel) * sizeof(float); +} + +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; +} + +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); + }); +} size_t vk_datatype_size(vkgraph::VkDataType dtype) { switch (dtype) { @@ -982,8 +1089,225 @@ void WebGPUGraph::build( tensors_[oid].nbytes}); } - // Phase 3: Build operator dispatch chain + 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; + + // Detect only the exact q4 gate/up MLP pattern. Full-chain occurrence and + // definition counts make every folded intermediate private to the pattern. const auto* chain = graph->chain(); + if (chain) { + 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 (value_types_[id] == ValueType::ValueList) { + for (int member : value_lists_[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); + } + } + + 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 || + value_types_[id] != 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) || tensor_mem_obj_ids_[gate_id] < 0) { + return; + } + + const auto& gate = tensors_[gate_id]; + const auto& up = tensors_[up_id]; + const auto& sigmoid = tensors_[sigmoid_id]; + const auto& silu = tensors_[silu_id]; + const auto& out = tensors_[out_id]; + if (!is_fp32_tensor(gate) || !is_fp32_tensor(up) || + !is_fp32_tensor(sigmoid) || !is_fp32_tensor(silu) || + !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) { + 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_fusion_ops.count(op) != 0; + } + if (overlaps) { + continue; + } + const size_t fusion_idx = swiglu_fusions.size(); + swiglu_fusions.push_back(fusion); + swiglu_gate_producers[fusion.gate_op] = fusion_idx; + swiglu_anchors[fusion.mul2_op] = fusion_idx; + swiglu_skipped_ops.insert(fusion.sigmoid_op); + swiglu_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 + // swiglu_skipped_ops entry. + for (unsigned op : pattern_ops) { + claimed_fusion_ops.insert(op); + } + } + } + + // Phase 3: Build operator dispatch 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 @@ -1042,6 +1366,11 @@ void WebGPUGraph::build( if (ops.size() != 3) { continue; // gate+up is a 2-group; o/down/lm_head are 1 each. } + if (std::any_of(ops.begin(), ops.end(), [&](unsigned op) { + return claimed_fusion_ops.count(op) != 0; + })) { + continue; + } // 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); @@ -1149,178 +1478,7 @@ void WebGPUGraph::build( qkv_member[ops[0]] = gidx; qkv_member[ops[1]] = gidx; qkv_member[ops[2]] = gidx; - } - } - } - - // 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}; + claimed_fusion_ops.insert(ops.begin(), ops.end()); } } } @@ -1342,34 +1500,28 @@ 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 SwiGluFusion& fusion = swiglu_fusions[anchor_it->second]; + 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; + } + + const size_t dispatch_begin = dispatches_.size(); // 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, @@ -1391,7 +1543,7 @@ void WebGPUGraph::build( create_scratch_buffer(tensors_[g.out_v].nbytes); } } - const size_t dispatch_begin = dispatches_.size(); + webgpu_operator_registry().get_op_fn(op_name)(*this, args); { @@ -1419,14 +1571,6 @@ void WebGPUGraph::build( add_qkv_fused_hook(qkv_groups[lit->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); - } - } const size_t dispatch_end = dispatches_.size(); @@ -1751,144 +1895,6 @@ void WebGPUGraph::add_qkv_fused_dispatch(QkvFusionGroup& g) { 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 diff --git a/backends/webgpu/runtime/WebGPUGraph.h b/backends/webgpu/runtime/WebGPUGraph.h index 9d0bebcea36..39911ff5ff1 100644 --- a/backends/webgpu/runtime/WebGPUGraph.h +++ b/backends/webgpu/runtime/WebGPUGraph.h @@ -595,15 +595,6 @@ class WebGPUGraph { // 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); }; } // namespace executorch::backends::webgpu 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/test/native/test_dynamic_shape.cpp b/backends/webgpu/test/native/test_dynamic_shape.cpp index ee821a6575e..971de490652 100644 --- a/backends/webgpu/test/native/test_dynamic_shape.cpp +++ b/backends/webgpu/test/native/test_dynamic_shape.cpp @@ -208,6 +208,97 @@ void run_qkv_routes(Module& m, int m_rows) { } } +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, kSwiGluK}, 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; @@ -799,6 +890,216 @@ 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"; + } + 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_qkv_fused"), 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/ops/dynamic_shape/test_dynamic_shape_export.py b/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py index 09dc03a3f0d..99e61c2efee 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 @@ -96,6 +96,78 @@ 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, + ) -> 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(64, output_width, bias=False).eval() + quantize_( + linear, + IntxWeightOnlyConfig(weight_dtype=torch.int4, granularity=PerGroup(32)), + ) + 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 + + 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.""" @@ -262,7 +334,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. @@ -286,6 +363,122 @@ def export_dynamic_shape_cases(out_dir: str) -> None: LIN_MAXM = 128 +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, SWIGLU_K)) + 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), + [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,