Repository navigation
[xegpu][examples] Remove replace_with_fused_attention and switch to dependent reduction fusion. - #298
Merged
tkarna merged 10 commits intoOct 12, 2026
Conversation
charithaintc
marked this pull request as draft
October 9, 2026 16:15
charithaintc
marked this pull request as ready for review
October 9, 2026 17:48
charithaintc
marked this pull request as draft
October 9, 2026 17:49
charithaintc
marked this pull request as ready for review
October 9, 2026 18:10
There was a problem hiding this comment.
🟡 Changes recommended
Fast-math annotations violate masked-infinity semantics, and the exported legacy API is removed without a compatibility period.
2 open findings
What changed in this PR
Replaces the custom fused-attention rewrite with dependent-reduction fusion across XeGPU schedules and payloads.
Changes:
- Adds shared online-attention fusion and fast-math transforms.
- Expresses softmax explicitly in attention payloads.
- Updates nanoGPT and Llama schedules, including causal masking.
| File | Description |
|---|---|
wg_tiling_schedule.py |
Uses shared attention fusion. |
vectorize_schedule.py |
Removes dead values after vectorization. |
fused_attention_schedule.py |
Adopts dependent-reduction fusion. |
attention_fusion.py |
Adds shared fusion utilities. |
gpu_attention_payload.py |
Emits explicit mixed-precision softmax. |
replace_with_fused_attention.py |
Removes the legacy transform. |
transform_ext/__init__.py |
Removes the legacy API export. |
nanoGPT_schedule.py |
Migrates attention scheduling. |
nanoGPT_payload.py |
Emits explicit softmax reductions. |
llama3.py |
Passes causal configuration into payload generation. |
llama3_schedule.py |
Migrates causal GQA fusion. |
llama3_payload.py |
Emits explicit causal attention dataflow. |
🧠 Review effort: Balanced
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+108
to
+115
| fast = transform.param_constant( | ||
| transform.AnyParamType.get(), ir.Attribute.parse("#arith.fastmath<fast>") | ||
| ) | ||
| for op_name in ("math.exp", "arith.subf"): | ||
| unmarked = structured.structured_match( | ||
| anytype, func, ops=[op_name], op_attrs={"fastmath": no_fastmath} | ||
| ) | ||
| transform.annotate(unmarked, "fastmath", param=fast) |
Contributor
Author
There was a problem hiding this comment.
Agree that fast is more extreme level of fastmath. In my understanding MathToXeVM depends on both having this flag. I would revisit this later to fine tune the fastmath flags required for this case.
Contributor
Author
There was a problem hiding this comment.
it turns out the subf does not need the flag here. so I removed it. thanks.
Comment on lines
19
to
21
| from .ops.trace_producers import trace_producers | ||
| from .ops.reverse_handles import reverse_handles | ||
| from .ops.update_address_space import update_address_space |
Contributor
Author
There was a problem hiding this comment.
renamed PR title to avoid the confusion.
tkarna
approved these changes
Oct 9, 2026
tkarna
left a comment
tkarna
left a comment
Contributor
There was a problem hiding this comment.
Very nice! Only minor comments.
charithaintc
changed the title
replace_with_fused_attention by using dependent reduction fusion. replace_with_fused_attention and switch to dependent reduction fusion.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR updates xegpu schedules with
fuse_dependent_reduction_opsto fuse attention-like workloads instead of relying on manualreplace_with_fused_attentionop.Note that
apply_attention_tilingalso uses the new approach.Depends on #297
This is a part of landing #274
Assisted by: Claude, Codex