Repository navigation
[XNNPACK] Delegate GEMMs that take parts of one packed weight or bias - #25005
john-rocky wants to merge 1 commit into
Conversation
get_static_param_deps now also accepts a weight or bias computation whose multi-output op, such as the split of nn.MultiheadAttention's packed in_proj_weight, gives its other outputs to GEMMs of the same op as the same argument. GEMMConfig.check_constraints accepts such a GEMM only together with those GEMMs, and only if their merged partition has no node with too many dims and creates no cycle. FoldComputedParamsPass folds the op inside the delegate without changes.
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/25005
Note: Links to docs will display an error until the docs builds have been completed.
|
meta-cla
Bot
added
the
CLA Signed
|
@JakeStevens here is the draft for the packed |
This PR needs a
|
Summary
Follows #24352, which leaves a computation shared by several GEMMs on portable. Now GEMMs that take the parts of one split, such as the packed projections of
nn.MultiheadAttention, are delegated together, and the split is folded inside their delegate.nn.MultiheadAttentionpacks the weights of its query, key and value projections into onein_proj_weightand their biases into onein_proj_bias; whenkdimorvdimdiffers from the embedding dim, the weights are separate params and only the biases are packed. Unless the query, key and value are one tensor, it splits each packed param withsplitorchunkand gives each part to its own linear. On main, the partitioner declines those linears because the split also feeds the other linears, so they run on the portable kernels in cross-attention, with different key and value tensors, withkdimandvdim, and innn.TransformerDecoderLayer.get_static_param_depsnow also accepts a weight or bias computation in which a multi-output op gives its other outputs to other GEMMs, as long as they are of the same kind as the GEMM being checked, such aslinear, and take their part in the same argument position. The config partitioner already merges partitions that share a node other than a constant, so these GEMMs become one partition, andFoldComputedParamsPass, which this PR does not change, folds the multi-output op into one constant per output.GEMMConfig.check_constraintsaccepts such a GEMM only together with the rest of its group. Every GEMM of the group must pass its own checks, and the merged partition of the group's GEMMs and their deps must have no node with too many dims and must passvalidate_partition, the cycle check offuse_as_graphmodule. The folded outputs together may be no larger than the params they free, as before.Without these checks, lowering fails. If one GEMM of a group is accepted and another is not, the one left out takes the multi-output op's outputs from the delegate, and
FoldComputedParamsPassraises. The same happens when a node with too many dims is in one GEMM's partition, becauseget_partitiondrops that GEMM aftercheck_constraintshas accepted it. A group with a cycle still becomes one partition, whichfuse_as_graphmodulerejects.#23154 folds constant
catandsplitchains before decomposition, intransform_for_pre_decomposition, and itstest_split_forms_delegate_and_match_compiled_runtimecovers three biases concatenated, split again and given to threelinearoraddmmcalls. This PR folds such a chain inside the delegate of the GEMMs that use it, the way #24352 folds a computed weight or bias. With that test'sProjectionBiasesmodule, all six cases lower to one delegate with no operator left outside it, and the outputs match eager.Lowered with
XnnpackPartitioner()on macOS arm64, embedding dim 16, 2 heads,batch_first=True, query(1, 5, 16), memory(1, 7, 16), default weights, biases drawn withtorch.nn.init.normal_, main → this PR:.ptebyteskdim=12, vdim=10nn.TransformerDecoderLayer,dim_feedforward=32The two extra delegate calls in
nn.TransformerDecoderLayercome from its key and value projection, which now shares a delegate with the query projection and so runs after self-attention. Thepermuteandcloneafter the key and value projection sit betweenunsqueeze,squeezeandselect, which stay on portable, so they no longer fit into other delegates and get one each.I have not measured latency.
This draft is for a look at the approach, as suggested on #24352. Suggested review order:
get_static_param_deps, thenGEMMConfig.check_constraintsand the two new methods behind it,_check_gemms_sharing_depsand_get_gemms_sharing_deps.Test plan
python -m pytest -v backends/xnnpack/test/passes/test_fold_computed_params.py16 passed, 26 subtests passed. The new tests lower linears and addmms that take the parts of a packed weight or bias; the parts are made with
split, withchunk, or withcatfollowed bychunk. They also lowernn.MultiheadAttentionin cross-attention, in cross-attention with dynamic query and memory lengths, with different key and value tensors, withkdimandvdim, and innn.TransformerDecoderLayer. In each of these,in_proj_biasis drawn at random, because it is zero by default. They check that no GEMM or split is left outside the delegates and that the runtime matches eager. For the linears and addmms they also check that a single delegate stores the parts, and not the packed param, as constants. A packed linear bias, and a packed linear weight with its bias, are also covered under static and dynamic quantization. These 14 subtests fail on main. The other new tests check what must not be delegated: a part also used by another op, a GEMM whose output passes through an op left on portable into the GEMM that takes the other part, parts that together are larger than the param, a GEMM of the group whose bias is computed with more dims than XNNPACK supports, and a conv1d that shares a split weight with a conv1d rejected for its dynamic shape. Each of the last four fails when its check is removed.python -m pytest -q backends/xnnpack/test/passes backends/xnnpack/test/ops backends/xnnpack/test/serialization backends/xnnpack/test/quantizer backends/xnnpack/test/test_xnnpack_partitioner.py -o addopts="" --continue-on-collection-errorsRun on macOS arm64 without a source build. The Python sources of this PR and of main each ran on top of the executorch 1.5.1 wheel, which supplies the native runtime and the pybindings:
7fa0b8a5b, the base of this PR): 69 failed, 518 passed, 1 skipped, 1 error, 195 subtests passedIn these runs, no failure is caused by this PR. Every test that fails on main also fails with this PR, which had one more failure. Most of the shared failures are
Failed to execute method forwarderrors from the 1.5.1 runtime, andtest_prelu.pyis not collected because it importsExecuTorchResult, which the 1.5.1 pybindings lack. The extra failure isTestLinear.test_qd8_per_channel_linear_parallel, which draws its inputs and batch size without a seed and compares withatol=0.05; the failing run had a max abs diff of 0.054. With seeds 0 to 29, main and this PR give the same max abs diff in every output comparison, and the same 4 seeds fail on both.Tests that fail on both
test_mean_dim.py: 22 teststest_io_dim_order.py: 7 teststest_upsample_nearest2d.py: 6 teststest_linear.py:test_linear_qd8_f16_per_channel_int4,test_linear_qd8_f16_per_token_weight_per_channel_group_int4,test_linear_qd8_f32_per_channel_int4test_gelu.py:test_fp32_gelu_approximation_2_tanh,test_fp32_gelu_approximation_3_tanhtest_bmm.py:test_bf16_bmmtest_softmax.py:test_bf16_softmaxtest_xnnpack_partitioner.py:test_enable_bf16_runs_an_attention_blocktest_xnnpack_quantizer.py:test_grutest_prelu.py: not collectedThis PR was authored with Claude Code.