Groverkss

How to write MLA as MHA

Terminology:

@ := Matrix Multiplication

Example: A @ B = matmul(A, B)

.T := Transpose

Groverkss

• Attention
  • Section 1 : What do attention variants really look like?
    • Variant class 1 : Score modifiers (Common)
    • Variant class 2 : Quantization and Not Really Attention? (Rare)
  • Section 2 : The Math Behind Attention Codegen
    • Reduction loops in attention
    • Online Softmax?
    • Attention with online softmax:
    • Online Softmax plus Matmul?

• Attention with Online Softmax+Matmul

bjacob

IREE/MLIR/Linalg tutorial

Introduction

This tutorial is simultaneously about IREE, MLIR, and specifically the MLIR Linalg dialect.

What is MLIR?

MLIR is a programming language, but MLIR in itself is almost just an empty shell. What it really provides is a framework allowing to define MLIR dialects which are where the features come from.

bjacob

Compile for LLVM-CPU, AArch64, with +i8mm extension, causing MaterializeEncoding to pick (8, 8, 8) tile sizes:

tools/iree-compile --iree-hal-target-backends=llvm-cpu --iree-llvm-target-triple=aarch64-none-linux-android29 --iree-llvm-target-cpu-features=+i8mm  --iree-flow-enable-data-tiling ~/matmul.mlir --mlir-disable-threading -o /tmp/a.vmfb --mlir-print-ir-after-all 2>/tmp/log

Compiler for VMVX with microkernels, causing dynamic tile sizes:

tools/iree-compile --iree-hal-target-backends=vmvx --iree-flow-enable-data-tiling --iree-vmvx-enable-microkernels ~/matmul.mlir --mlir-disable-threading -o /tmp/a.vmfb --mlir-print-ir-after-all 2>/tmp/log