Fix GQA fusion to produce present key/value (#2634)

justinchuby · web-flow · commit 75b3d42e3416 · 2025-10-15T11:01:40.000-07:00
Output present key value from the Attention op because past key value is provided. Previously the Attention op created would consume past key/value but not produce present key/value, which is not correct for ORT. <img width="1377" height="1225" alt="image" src="https://github.com/user-attachments/assets/118958b4-bc27-4912-b70b-000549887c0f" /> Replaces #2632 Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
diff --git a/onnxscript/rewriter/rules/fusion/_gqa.py b/onnxscript/rewriter/rules/fusion/_gqa.py
@@ -52,7 +52,7 @@ def pattern(
             _outputs=["attention_BHSDh"],
         )
 
-        return attention_BHSDh
+        return attention_BHSDh, present_key_BHkvStD, present_value_BHkvStD
 
     def check(
         self,
@@ -103,6 +103,7 @@ def rewrite(
             past_key_BHkvSpD,
             past_value_BHkvSpD,
             **original_attrs,
+            _outputs=3,
         )
 
 

Original file line number	Diff line number	Diff line change
`@@ -52,7 +52,7 @@ def pattern(`
`52`	`52`	`_outputs=["attention_BHSDh"],`
`53`	`53`	`)`
`54`	`54`
`55`		`- return attention_BHSDh`
	`55`	`+ return attention_BHSDh, present_key_BHkvStD, present_value_BHkvStD`
`56`	`56`
`57`	`57`	`def check(`
`58`	`58`	`self,`
`@@ -103,6 +103,7 @@ def rewrite(`
`103`	`103`	`past_key_BHkvSpD,`
`104`	`104`	`past_value_BHkvSpD,`
`105`	`105`	`**original_attrs,`
	`106`	`+ _outputs=3,`
`106`	`107`	`)`
`107`	`108`
`108`	`109`