[dependencies] Install torch-mlir from PyPI by tkarna · Pull Request #302 · llvm/lighthouse

tkarna · 2026-10-09T21:42:46Z

torch-mlir is now installed from PyPI.
Bumps torch-mlir version to 20261001.
Bumps mlir-python-bindings version.
Imported IR is now slightly different for some kernels, various fixes in lowering schedules.

Addresses #283.

Assisted-by: Claude, GPT

tkarna · 2026-10-12T09:04:43Z

torch-mlir behavior has changed in the latest version:

Handling of empty-sized tensors seems broken
- Added a workaround in examples/llama/ref_model.py
scaled_dot_product_attention lowering has changed
- It now includes a boolean mask that masks out possible -inf values in the score matrix
  - xegpu lowering was updated to handle such mask in the graph.
- Most of the computation is now done in wide datatype f32 instead of bf16. Consequently, the gemms are now (f32, f32)->f32.
  - Added narrow_vector_contract_operands.py to cast vector.contract ops to DPAS-compatible (f16, f16)->f32 form.
- The graph has new elementwise producer ops before gemms/reduction
  - xegpu lowering was updated to correctly tile and fuse them.

tkarna · 2026-10-12T09:06:12Z

FYI @charithaintc @nbpatel

tkarna · 2026-10-14T20:40:53Z

Rebased on current main pulling in #298. Making the new transform_to_flash_attention method to work with the -inf masking (pytorch's safe softmax) required quite substantial changes.

Most importantly, in the attention lowering we can no longer call the linalg-fuse-elementwise-ops pass because it will fuse the inf mask into the PV contraction op forming a 3D contraction op. Omitting this pass, we not only lack the elemenwise fusion but also many other cleanup patterns that this pass applies. Therefore

We now apply linalg.erase_unused_operands_and_results pattern that was recently exposed as a transform pattern. Requires bumping Lighthouse llvm version.
linalg::populateFoldReshapeOpsByExpansionPatterns lacks transform dialect bindings. Added a temporary fold_reshape_ops_by_expansion transform op as a workaround.
FoldScalarOrSplatConstant pattern is also missing. Added a temporary fold_scalar_or_splat_constant transform op as a workaround.

Without the elementwise fusion the attention graph is now more complex with smaller ops. Several generalizations were implemented in sink_normalization_past_contraction and fuse_dependent_reduction_ops methods to support the new graph form.

The fastmath flag was also changed from fast to afn in the attention layer to work correctly in the presence of -inf values.

charithaintc

added some comments to clear some initial doubts.

charithaintc · 2026-10-14T22:52:27Z

-    result, and R2 reduces E's result. The loop step supplies the tile size.
+    result, and R2 reduces E's result, possibly through elementwise intermediates.
+    The loop step supplies the tile size. Intermediates are inlined into a
+    temporary consumer for analysis only. In the fused loop they are tiled as


is this temp consumer dead code? is that removed after analysis?

I think proper way to handle additional elementwise ops would be to fuse them into a single E. did you consider using fine grained elementwise fusion extension for this?

is this temp consumer dead code? is that removed after analysis?

Yes. We create a temporary fused op to allow easier analysis of the graph, and it is being explicitly removed. The benefit is that we can reuse the existing legality checks without a need to write new analysis methods for the non-fused case (e.g. proving linearity or tracing indexing maps).

I think proper way to handle additional elementwise ops would be to fuse them into a single E. did you consider using fine grained elementwise fusion extension for this?

The caveat here is the -inf masking. We actually do not want to fuse the arith.select op into E. Therefore, in case of a masked op, we always have a graph E -> zero-select -> PV and handling this pattern would require more complex analysis in any case. That said, there certainly are different ways of implementing this.

I see. thanks!

charithaintc · 2026-10-14T22:54:34Z

-    """Create the fusion transform and return its new loop handle."""
+    """Fuse the chain using apply-time payload analysis.
+
+    Recognized softmax uses stable correction, assumes finite auxiliary inputs


In follow up PR, we should try to remove this specialization for softmax. I think this ok for now given the complexity of the changes.

Yes, this transform now has softmax specific handling. We essentially need two guards in the fused loop:

In the E op we need to set exp(-inf - (-inf)) to 0 to guard against producing NaNs.

Replace the correction term term(a) / term(b) with stable a == -inf ? 0 : exp(a + b). I.e. two guards, first against -inf and second one for exp(a+b) stability.

These are not generic guards but depend on E being a softmax. I agree this is not ideal. In principle it could be possible to pull these guards out of the fuse_dependent_reduction_ops transform into a separate softmax specific transform that runs after. It would be somewhat complex because we'd need to trace the IR, check for softmax and match the appropriate handles etc.

I also note that convert-math-to-xevm pass converts exp(a)/exp(b) to exp(a-b) later on in the pipeline so for xegpu this guard is not needed here.

I did experiment with dropping these guards and the KB attention layer kernel executes correctly. So one options is to drop these in this PR and handle them later.

I also note that convert-math-to-xevm pass converts exp(a)/exp(b) to exp(a-b) later on in the pipeline so for xegpu this guard is not needed here.

OK. I think I have a better understanding now.

As you mentioned already we reply on fastmath flags to do the reassociation. Without fastmath this is not a valid rewrite ( exp(a)/exp(b) to exp(a-b)). Honestly I don't think we need to guard against it. user needs to add reassoc flag and legality check should check them before fusion in my view.

charithaintc · 2026-10-14T23:14:41Z

+class RewriteAttentionRowMaskOp(
+    TransformExtensionDialect.Operation, name="rewrite_attention_row_mask"
+):
+    """Replace all--inf row masks using guarded online attention's maximum.


can you explain how this works?

several questions:

does this run after or before dependent red fusion? if after fusion then how is the AND reduction fused>

does this remove the AND reduction?

Yes, this transform runs after the row-sum and PV fuse_dependent_reduction_ops transforms. It's purpose is to fuse the scores == -inf mask computation to the reduction loop by re-using the already existing max reduction (note that all(scores == -inf) is equivalent to max(scores) == -inf):

In the reduction loop, check whether the running max reduction is -inf.

After the reduction loop, check whether the final maximum is -inf. With this update we do not need the and reduction anymore and it can be dropped.

This is a specialized transform. The benefit is that it's kept separate from the more generic fuse_dependent_reduction_ops op.

okay. it makes sense.

for future: Do you think there will be a way to fuse the scores == -inf similar to how we fuse others here. maybe converting it to a OR reduction of score != -inf.

That could be a possible way to handle and reductions, let's assess these ideas once this PR has landed.

charithaintc · 2026-10-14T23:22:06Z

 sum accumulator is ``exp(-m_new) / exp(-m_old)`` before adding this contribution
 of current tile ``sum(exp(x_tile - m_new))``.

+Recognized softmax chains use ``exp(m_old - m_new)`` and zero weights for


it seems this removes the dependency of row mask on PV? and moves it after the fused loop.

it also looks like code is specialized for the softmax case now. this is less desirable. but I guess we can think of a proper solution in another PR.

I also wonder if we can do the RewriteAttentionRowMaskOp can be applied before fusion so that PV does not directly depend on the mask output. instead we compute the mask normally and do mask application after the fused loop.

I also wonder if we can do the RewriteAttentionRowMaskOp can be applied before fusion so that PV does not directly depend on the mask output. instead we compute the mask normally and do mask application after the fused loop.

This kind of design might be possible, but simply computing the mask from the final max before dependent-reduction fusion is insufficient. The in-loop PV computation needs a predicate based on the updated running max, so fusion would need to remap that dependency. Also, moving the mask application after PV is not generally equivalent: an all--inf intermediate prefix can produce NaNs in the running computation even when the final row contains finite scores.

okay. I am still missing something here.

all(scores == -inf) whole point of this is to avoid that row in P participating in PV, right? otherwise things would underflow.

My suggestion was to not use running max at all. instead we do the all(scores == -inf) reduction normally inside the fused loop but fused PV does not use it. PV will possibly will overflow in certain rows.

Then after the loop we rescale PV with the output of all(scores == -inf). this will mask out rows of PV where P had all -inf row. Maybe my understanding is not correct. but good to verify if this approach works or now.

(btw even if this works we don't have a general way to fuse the AND reduction in first place. so its an open problem)

This is slightly more subtle. On the global untiled kernel level, we indeed only need to compute the full -inf rows and we can apply that mask after PV matmul to mask out those invalid values.

However, this changes when we apply online reduction to the softmax. Now, if the first n values in that row are -inf, the partial max is -inf and softmax gives exp(-inf - -inf) = NaN. This NaN persists and propagates to the output, even if the later values were finite. Thus opting for online reduction exposes a new NaN generation mechanism(!). And if we want to guard against these NaNs as well, we need to update the mask computation accordingly, i.e. update it similarly with online reduction.

That's why we have both rewrite_attention_row_mask transform (in the reduction loop, it uses partial max_reduction to mask the p tile, and after the loop also applies the mask to the final result) and the softmax-specific guarding in fuse_dependent_reduction_ops.

ok. understood.

charithaintc

LGTM. may require some clenups/redesigns in follow up PRs.

thanks for clearing the misunderstandings.

This is required with attn layer in newer torch-mlir version 20261001.

If attention layer has a P row mask, we need to rewrite the mask usage in PV contraction to be able to fuse the mask generation in the reduction loop.

…ms referenced by their inputs This allows CSE to eliminate identical factors in the graph.

…aling loop

This should be replaced by adding bindings to linalg::populateFoldReshapeOpsByExpansionPatterns.

We cannot call linalg-fuse-elementwise-ops pass because the p mask must be kept separate from the pv matmul to be able to map the latter to a xegpu.dpas op. Skipping the monolithic fusion, the graph now has more small ops and lowering logic needs to be updated. - fuse_dependent_reduction_ops.py: _prepare_consumer still builds an inlined copy of R2, but only to run the legality and correction-factor checks. It then deletes that copy and passes the real R2 plus the list of in-between elementwise ops to the loop builder. Keeps the mask separate from the pv matmul. - dependent_reduction_fusion.py: _chain_tiled_dims works out which loop dims of each in-between op follows R2's reduction axis, and rejects the fusion if that isn't a plain dimension. The in-between ops are cloned in reduction loop and R2 reads the result of the last op. - attention_row_mask.py: handles non-fused attention IR, e.g. the mask select and -inf compare as separate ops, zero as a dense constant tensor.

This should be replaced by upstream FoldScalarOrSplatConstant pattern.

…o linalg

Also move hoisting to a better position

…n-tiled level

…erns

Necessary so that masked attn can be lowered to binary

torch-mlir now keeps the intermediate values at f32 so we have a truncf op between xegpu load and dpas

tkarna marked this pull request as ready for review October 12, 2026 08:52

tkarna requested review from adam-smnk and rengolin October 12, 2026 09:04

rengolin approved these changes Oct 12, 2026

View reviewed changes

adam-smnk approved these changes Oct 12, 2026

View reviewed changes

tkarna force-pushed the pypi-torch-mlir branch from ce0911b to 2a2e567 Compare October 14, 2026 20:14

charithaintc reviewed Oct 14, 2026

View reviewed changes

charithaintc approved these changes Oct 15, 2026

View reviewed changes

tkarna added 19 commits October 16, 2026 15:49

install torch-mlir from PyPI and pin latest version

faa5fb4

sink_normalization_past_contraction: support select mask in op chain

7158313

This is required with attn layer in newer torch-mlir version 20261001.

fuse_dependent_reduction_ops supports P inf masking

d3b074e

transform_ext: add rewrite_attention_row_mask

a3e9915

If attention layer has a P row mask, we need to rewrite the mask usage in PV contraction to be able to fuse the mask generation in the reduction loop.

fuse_dependent_reduction_ops: compute correction factors only over di…

51eed19

…ms referenced by their inputs This allows CSE to eliminate identical factors in the graph.

fused_attention_schedule: update wg tiling matchers

0f6635e

attention_fusion: update matching and transforms to work with P mask

e5baed6

sink_normalization_past_contraction: generalize to handle separate sc…

c8d447b

…aling loop

transform_ext: add fold_reshape_ops_by_expansion op

fb36c62

This should be replaced by adding bindings to linalg::populateFoldReshapeOpsByExpansionPatterns.

attention_fusion: drop linalg-fuse-elementwise-ops pass

00ca284

fused_attention_schedule: use fold_reshape_ops_by_expansion ops

0de9aff

attention_fusion: fuse P producers into reduction loop

e288bfa

annotate_fastmath_flags: use more conservative afn flag

02db9b0

attention_fusion: simplify final clean up

d4b2ae9

transform_ext: add fold_scalar_or_splat_constant op

b37d59c

This should be replaced by upstream FoldScalarOrSplatConstant pattern.

fused_attention_schedule: ise fold_scalar_or_splat_constant op

dac6157

transform_ext: add narrow_vector_contract_operands op

db24299

lowering_common: vectorize applies vector.contract operand narrowing

a0cfe18

tkarna added 17 commits October 16, 2026 15:50

kernel_bench: force torch-mlir SDPB math backend to enable lowering t…

56fb538

…o linalg

kernel_bench: tweak attention q load tile to enable lowering

ca282b7

pyproject.toml: update llvm version

f8f5307

fused_attention_schedule: tweak comment

d9dd187

attention_fusion: fix producer fusion for non-masked case

68e7f4a

fused_attention_schedule cleanup

180bace

lowering_common: vectorize includes rm-dead-code pass

950361f

Also move hoisting to a better position

fused_attention_schedule: fold constants and add fastmath in reductio…

dc8e2aa

…n-tiled level

cleanup_schedule: replace linalg elemwise fuse pass with cleanup patt…

6c0986d

…erns

wg_tiling_schedule: update to work with new flash attn transform

b394c79

vectorize_schedule: remove-dead-values pass is in vector helper now

cd7c356

annotate_layouts_schedule: tweak attention q load tile size

f8de8fb

Necessary so that masked attn can be lowered to binary

mlp_schedule: fix gemm k tiling

55606d7

wg_tiling_schedule: fixes for gemm/reduction/attention cases

7cabed0

fix torch llama example

445f816

mlp_schedule: fix MLP dpas annotation

70e9678

torch-mlir now keeps the intermediate values at f32 so we have a truncf op between xegpu load and dpas

annotate_layouts_schedule: fix MLP dpas layout annotation

6d06898

tkarna force-pushed the pypi-torch-mlir branch from e4501d2 to 6d06898 Compare October 16, 2026 12:51

tkarna merged commit e207770 into llvm:main Oct 16, 2026
3 checks passed

tkarna deleted the pypi-torch-mlir branch October 16, 2026 13:20

Conversation

tkarna commented Oct 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

tkarna commented Oct 12, 2026

Uh oh!

tkarna commented Oct 12, 2026

Uh oh!

tkarna commented Oct 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

charithaintc left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tkarna Oct 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

charithaintc left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

tkarna commented Oct 9, 2026 •

edited

Loading

tkarna commented Oct 14, 2026 •

edited

Loading

tkarna Oct 15, 2026 •

edited

Loading