Skip to content

[xegpu] Add torch llama3 attention example - #307

Open
tkarna wants to merge 15 commits into
llvm:mainfrom
tkarna:xegpu-torch-llama3-attention
Open

tkarna wants to merge 15 commits into
llvm:mainfrom
tkarna:xegpu-torch-llama3-attention

Conversation

@tkarna

@tkarna tkarna commented Oct 16, 2026 •

Copy link
Copy Markdown
Contributor
  • Adds a KernelBench-like llama-bench_1B_attention.py kernel file which defines a Grouped-Query Attention (GQA) kernel following Llama-3.2-1B model specification.
  • The kernel is executed with kernel_bench.py routine via the new -i/--input flag.
  • The attention lowering schedules are updated to correctly lower the GQA layer. More IR canonicalization is required to fold 4D/5D expand/collapse_shape ops.
  • In transform_to_flashattention the reduction_tile argument is now a (zero-padded) vector to allow easier handling of inferred tile sizes.
  • Adds expand_output_destination transform op which rewrites materialize_in_destination(collapse_shape(X), D) into materialize_in_destination(X, memref.expand_shape(D)). This is required to avoid materializing a copy buffer for the layer result.
  • Minor fix in contraction_normalization.

@tkarna
tkarna force-pushed the xegpu-torch-llama3-attention branch from a3a027a to c05a9cb Compare October 16, 2026 21:11
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant