fix(functions): add optional bias slot to PackedMultiHeadAttention fallback - #449
Conversation
…llback The ORT com.microsoft::PackedMultiHeadAttention op has a 6-7 input positional signature: query, key(opt), value(opt), bias(opt), token_offset, cumulative_sequence_length, attention_bias(opt). Because token_offset/cumulative_sequence_length occupy positional slots 5 and 6, the optional bias at slot 4 must exist as a formal input. The fallback ir.Function declared only 5 formals (dropping the bias slot), so onnx-genai's function-inline admission rejected the 6-input call sites with a FunctionArityMismatch (call.input.len()=6 > func.input.len()=5) for the Qwen3.6-35B-A3B vision_encoder. Insert the bias formal between value and token_offset. The body ignores bias (the block-diagonal bias is reconstructed from cu_seqlens); it is present only to preserve positional slot alignment. Add a regression test asserting the 6 formals and their positional order. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Performance Comparison
|
There was a problem hiding this comment.
Pull request overview
This PR fixes the standard-ONNX fallback ir.Function signature for com.microsoft::PackedMultiHeadAttention by adding the missing optional bias input slot so function inlining/admission no longer fails due to positional arity mismatch.
Changes:
- Add
biasas the 4th formal input (slot-alignment fix) in the PackedMHA fallback function body. - Update the fallback function docstrings/comments to document the unused-but-required
biasslot. - Add a unit test asserting the fallback function’s formal input order and identity.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
src/mobius/functions/packed_multi_head_attention.py |
Adds the missing bias formal input (slot 4) to match the PackedMHA positional signature and updates docs/comments accordingly. |
tests/packed_multi_head_attention_function_test.py |
Adds regression tests to lock down the fallback function’s formal-input order/arity and identity. |
| """Unit tests for the PackedMultiHeadAttention fallback ``ir.Function``. | ||
|
|
||
| These guard the formal-input arity/order of the standard-ONNX fallback for | ||
| ``com.microsoft::PackedMultiHeadAttention``. ORT's op has a 6-7 input |
|
VERDICT: APPROVE Independent review by Harry (reviewer, opus-4.8). I did not author this change. Reviewed at commit 864c04f on branch squad/packed-mha-bias-slot in a detached worktree off origin/main. 1. Positional-wiring integrity — CORRECT. 2. bias is genuinely inert. 3. Call-site parity — CORRECT (all 3 sites). Note: the actual call sites are
4. Test non-vacuous + mutation check. 5. Scope. Test result: The fix correctly restores the 6th formal (bias at slot 4) so onnx-genai's function-inline admission accepts the 6-input call nodes, with zero behavioral change to the fallback decomposition. Approving. |
|
@copilot please fix review comments |
… test Rename the computed attention bias in the module docstring pseudo-code to attn_bias so it is not confused with the new formal `bias` input, and move the fallback-function test next to its implementation under src/mobius/functions/ to match the existing convention.
|
|
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 2 out of 2 changed files in this pull request and generated no new comments.
Suppressed comments (1)
src/mobius/functions/packed_multi_head_attention_test.py:55
test_packed_mha_admits_six_input_callis currently a tautology: it only checks6 <= len(func.inputs), which is already guaranteed bytest_packed_mha_declares_six_positional_formal_inputs. This doesn’t actually exercise the failure mode described in the docstring (InlinePass/function admission rejecting a 6-input call node). Consider rewriting this test to build a minimal model containing acom.microsoft::PackedMultiHeadAttentionnode with a missing bias slot (""), register the function body, and runInlinePassto ensure inlining succeeds and the node is expanded.
def test_packed_mha_admits_six_input_call() -> None:
func = packed_multi_head_attention()
# A call site emits 6 inputs: (q, k, v, "", token_offset, cu_seqlens).
# onnx-genai admits a call when len(call inputs) <= len(func inputs), so
# the function must declare at least 6 formals for the call to be inlined.
call_input_count = 6
assert call_input_count <= len(func.inputs)
Summary
Fixes the formal-input arity of the standard-ONNX fallback
ir.Functionforcom.microsoft::PackedMultiHeadAttentionby adding the optionalbiasslot.Root cause
ORT's
PackedMultiHeadAttentionhas a 6–7 input positional signature:Because
token_offset/cumulative_sequence_lengthoccupy positional slots 5 and 6, the optionalbiasat slot 4 must exist as a formal input (present-as-absent at call sites via"").Call sites in mobius already correctly emit 6 inputs
(q, k, v, "", token_offset, cu_seqlens):src/mobius/models/_qwen3_vl_vision.pysrc/mobius/rewrite_rules/_packed_attention.pyBut the fallback
ir.Functiondeclared only 5 formals, dropping thebiasslot. This made onnx-genai's function-inline admission reject the call:for the Qwen3.6-35B-A3B
vision_encoder. onnx-genai is spec-correct here; the bug was purely the missingbiasformal in mobius.Fix
biasas the 4th formal input, betweenvalueandtoken_offset, yielding the correct positional order:query, key, value, bias, token_offset, cumulative_sequence_length(6 formals).bias_inputparameter to the bodybody(...)in the same position. The body does not usebias— the fallback reconstructs the block-diagonal attention bias fromcumulative_sequence_length— so it is present only to preserve positional slot alignment (treated like the already-unusedtoken_offset_input).bias: (optional) — unused in fallbackin the correct positional order.Test
Added
tests/packed_multi_head_attention_function_test.py, which asserts:query, key, value, bias, token_offset, cumulative_sequence_length(not just the count —biasmust be at index 3);len(call inputs) <= len(func inputs));com.microsoft/PackedMultiHeadAttention).All 3 new tests pass;
tests/ep_optimization_test.py(30 tests) passes as a regression check; ruff check + format clean.