[xegpu] Generalize fused attention schedule for KernelBench by tkarna · Pull Request #266 · llvm/lighthouse

tkarna · 2026-08-24T16:41:32Z

Generalizes fused_attention_schedule so that it works for KernelBench level1-97 ScaledDotProductAttention kernel.

The schedule can consume 4D attention kernels [batch_size, n_head, n_ctx, d_head] as well as versions where the 2 leading dimensions have been collapsed to a single batch dimension batch_size, n_head within the payload function (e.g., in the examples/xegpu/fused_attention.py).

At the moment, the KernelBench problem shape must be simplified from the original [32, 32, 512, 1024] to [32, 32, 512, 64] because we do not yet have a good strategy for distributing the last dimension (head dimension).

fused_attention_schedule uses common xegpu helper utils where possible.
Updates lowering_common and inspect_kernel methods.
replace_with_fused_attention transform ext op is generalized to handle bf16/mixed data types and 2 leading dimensions.
The attention layer related docs and variable names are unified across examples/schedules/payload generation. For example, now the dimensions are called [batch_size, n_head, n_ctx, d_head] consistently everywhere.
examples/xegpu/kernel_bench.py supports level1-97 attention kernel. Currently the kernel file must be manually patched to set embedding_dimension = 64.

Assisted-by: Claude

tkarna · 2026-08-24T16:41:51Z

FYI @charithaintc

Copilot

Pull request overview

Generalizes the XeGPU fused-attention lowering pipeline to support KernelBench’s level1-97 ScaledDotProductAttention-style kernels, including support for both true 4D attention shapes and payloads that collapse the leading batch dimensions.

Changes:

Refactors common lowering utilities (get_payload_func, vectorization/bufferization/GPU launch conversion helpers) and updates existing schedules to use them.
Generalizes fused attention scheduling and the replace_with_fused_attention transform to handle additional shapes and mixed/bf16-related type flows.
Extends KernelBench example driver to recognize attention-like kernels and route them through the fused-attention schedule.

Reviewed changes

Copilot reviewed 11 out of 11 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
lighthouse/utils/mlir.py	Extends layer metadata extraction for batch_matmul to include element-type info.
lighthouse/schedule/xegpu/reduction_schedule.py	Updates schedule to use `get_payload_func` and updated helper signatures.
lighthouse/schedule/xegpu/mlp_schedule.py	Updates schedule to use `get_payload_func`.
lighthouse/schedule/xegpu/lowering_common.py	Introduces `get_payload_func` and refactors helper APIs; adds vector cleanup patterns and LICM pass in vector→xegpu.
lighthouse/schedule/xegpu/fused_attention_schedule.py	Reworks fused attention schedule to support more shapes and reuse common lowering helpers.
lighthouse/schedule/xegpu/elemwise_schedule.py	Updates schedule to use `get_payload_func`.
lighthouse/ingress/mlir_gen/gpu_attention_payload.py	Unifies attention dimension naming and updates payload generation for collapsed-batch handling.
lighthouse/dialects/transform/transform_ext/ops/replace_with_fused_attention.py	Generalizes fused-attention replacement for batch dims and mixed/bf16-related typing behavior.
examples/xegpu/nanoGPT_schedule.py	Updates callsite to updated `convert_to_gpu_launch` signature.
examples/xegpu/kernel_bench.py	Adds attention schedule path, adds CPU reference option, and adjusts verification tolerances.
examples/xegpu/fused_attention.py	Unifies naming, extends CLI params, and updates schedule parameter plumbing.

Suppressed comments (1)

lighthouse/dialects/transform/transform_ext/ops/replace_with_fused_attention.py:516

Call site uses the misspelled helper name normalize_ouput_by_sum; rename it to match the corrected helper name.

                output_normalized = normalize_ouput_by_sum(
                    pv_out, l_i_out, batch_shape, wg_rows, d_head, compute_type
                )

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

… dims

charithaintc

Looks great. some nit comments.

I will continue the review tomorrow.

charithaintc · 2026-08-24T23:23:00Z

                        compute_type,
                    )
+                    # Truncate Q@K^T (f32 accumulator) to the softmax type before scaling.
+                    if reduction_type != compute_type:


line 364 sets the reduction_type = compute_type. this this branch is never taken?

You are right, it's actually dead code. I'm inclined to keep the code as it is, however, because one can change the reduction data type to a narrower one just by changing the assignment reduction_type = compute_type. The generated IR is still valid. Added a comment.

charithaintc · 2026-08-24T23:23:12Z

+                    # Rescale running P@V accumulator by alpha; the accumulator is
+                    # kept in f32, so widen the softmax-type alpha to match.
+                    alpha_wide = alpha
+                    if reduction_type != compute_type:


charithaintc · 2026-08-24T23:23:32Z

            with ir.InsertionPoint.after(loop):
+                # The sum accumulator is in the softmax type; widen it to the f32
+                # accumulator type before dividing the f32 P@V result.
+                if reduction_type != compute_type:


charithaintc · 2026-08-24T23:25:38Z

-    - Q @ K^T:        2 * n_ctx^2 * n_head FLOPs
-    - Attention @ V:  2 * n_ctx^2 * n_head FLOPs
-    The softmax is left out: it is O(n_ctx^2) (~2% of the above at n_head = 64) and
+    - Q @ K^T:        2 * n_ctx^2 * d_head FLOPs


nit: n_head -> d_head change also need to be reflected on nanoGPT case

Updated all nanoGPT files to use the same variable naming convention

charithaintc · 2026-08-24T23:26:17Z


    # Check values match reference
-    values_ok = np.allclose(output, output_ref, rtol=1e-3, atol=1e-3)
+    values_ok = np.allclose(output, output_ref, rtol=1e-3, atol=5e-3)


what is the reason for atol change?

After updating llvm version the tests no longer passed, even with LH main. After scratching my head for a while I realized it's just a tolerance issue (test failed even with older llvm version if inputs were filled differently). So I think the tolerance was too tight. This kernel has two matmuls and a reduction so errors accumulate.

charithaintc · 2026-08-24T23:28:38Z

+    # loads, but are distributed over the subgroups.
+    prefetch_sg_data = layer_params["prefetch_tile"]
+    prefetch_sg_layout = [
+        reduction_tile // prefetch_sg_data[0],


init: maybe validate against num_sgs?

hmm, in principle you can use different number of sgs for the computation and prefetches? Might not be optimal config but who knows. The hardware limits should not be exceeded but those are not validated in the schedule.

charithaintc · 2026-08-24T23:31:09Z

+            op_attrs=op_attrs,
+        )
+    # Return the first function
+    func = transform_ext.extract_handle(func_ops, 0)


this would fail silently if none matched. what is the reason for removing match_and_split?

extract_handle raises an IndexError if the index is out of bounds so it's a definitive failure at lowering time.

also: - C -> n_embd - hidden -> n_ffn - vocab -> n_vocab - clean up unused variables

charithaintc

LGTM. thanks!

charithaintc · 2026-08-26T16:40:24Z

-            result_ref = torch_model(*torch_inputs).to("cpu")
+        else:
+            # execute torch model on the device
+            torch_inputs = [inp.to("xpu") for inp in torch_inputs]


is the recent results discrepancy only for GPU?

yes, this script only targets xegpu

hmm, I mean I usually cross check xegpu output with CPU output. it seems here you use pytoch with xpu support to cross validate?

Ah yes, in this script we indeed run the reference solution on the GPU with torch xpu backend. The new attention kernel is an exception though because torch xpu doesn't currently support torch's scaled_dot_product_attention op. That's why I added the --compute-reference-on-cpu flag to enable testing against torch CPU backend.

tkarna requested review from adam-smnk and a lite review from Copilot August 24, 2026 16:41

Copilot started reviewing on behalf of tkarna August 24, 2026 16:42 View session

tkarna force-pushed the xegpu-kb-attention branch from 1a63c1b to dc69604 Compare August 24, 2026 16:44

Copilot AI reviewed Aug 24, 2026

View reviewed changes

tkarna added 8 commits August 24, 2026 20:18

replace_with_fused_attention: support bf16/mixed dtypes and 2 leading…

a9c6ceb

… dims

lowering common: include transforms required by attention payload

37b0f55

lowering common: use more generic get_payload_func helper

73e7c3b

inspect payload: return batch_matmul element types

780f04e

fused attention schedule: revise and generalize

a6a3cc4

update fused_attention example

e63b750

kernel bench: support l1 attn benchmark

797f509

fused attention example: loosen error tolerance

7a92481

tkarna force-pushed the xegpu-kb-attention branch from dc69604 to 7a92481 Compare August 24, 2026 17:19

address comments

4436f21

charithaintc reviewed Aug 24, 2026

View reviewed changes

tkarna added 2 commits August 25, 2026 23:30

replace_with_fused_attention: add comment about reduction data type

501fefc

nanoGPT: use the same batch_size, n_head, n_ctx, d_head variable names

e467ca8

also: - C -> n_embd - hidden -> n_ffn - vocab -> n_vocab - clean up unused variables

charithaintc approved these changes Aug 26, 2026

View reviewed changes

tkarna merged commit 8fb00f7 into llvm:main Aug 26, 2026
3 checks passed

tkarna deleted the xegpu-kb-attention branch August 26, 2026 19:11

Uh oh!

Conversation

tkarna commented Aug 24, 2026

Uh oh!

tkarna commented Aug 24, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

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!

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 Aug 25, 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!

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!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

tkarna Aug 25, 2026 •

edited

Loading