Skip to content

[Performance] Fix K/V prefetch offsets and register spilling in flash attention tests #1226

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
charithaintc merged 1 commit into from
Aug 19, 2026
Merged

[Performance] Fix K/V prefetch offsets and register spilling in flash attention tests #1226

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 15 additions & 13 deletions test/Integration/Dialect/XeGPU/SG/flash_attention_fwd.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
// RUN: imex-opt %s --gpu-lower-to-xevm-pipeline="xegpu-op-level=subgroup" \
// RUN: imex-opt %s --gpu-lower-to-xevm-pipeline="xegpu-op-level=subgroup igc-cmd-options=-ze-opt-large-register-file" \
// RUN: | mlir-runner \
// RUN: --shared-libs=%mlir_levelzero_runtime \
// RUN: --shared-libs=%mlir_runner_utils \
Expand All @@ -12,6 +12,10 @@
// --runner mlir-runner -e main \
// --entry-point-result=void \
// --shared-libs=%irunner_utils,%mlir_runner_utils,%mlir_c_runner_utils,%mlir_levelzero_runtime --filecheck

// -ze-opt-large-register-file gives the kernel 256 GRF. At the default 128 GRF
// the loop-carried accumulators do not fit and IGC spills (~7.4 KB, ~234
// spill/fill accesses in the loop body).
module @flash_attention attributes {gpu.container_module} {
gpu.module @flash_attention_fwd {
gpu.func @flash_attention_fwd(
Expand Down Expand Up @@ -74,16 +78,14 @@ module @flash_attention attributes {gpu.container_module} {
%v_tile_slice = xegpu.create_nd_tdesc %V, shape: [%size_x, %BLOCK_DMODEL], strides: [%BLOCK_DMODEL, %c1] : memref<?x?xf16> -> !xegpu.tensor_desc<16x16xf16>

// K prefetch.
// Prefetch 16x32 tiles in 4x2 layout to cover 64x64
// x offset for prefetch is same as for q tiles. This means that WGs assigned to same batch also collaborate on prefetching
// the K, V tiles.
// NOTE: We also tried WGs prefetching from the begining of the K, V tiles but that did not work well because multiple
// WGs compete to prefetch the same data.
// Prefetch 16x32 tiles in 4x2 layout to cover 64x64.
// x offset for prefetch is the same as for the K, V loads (%wg_x_offset),
// so the prefetch tracks the loads BLOCK_N_3 rows ahead.
%sg_layout_x = arith.divui %sg_id, %c2 : index
%sg_layout_y = arith.remui %sg_id, %c2 : index

%prefetch_offset_x_t0 = arith.muli %sg_layout_x, %c16 : index
%prefetch_offset_x = arith.addi %wg_q_x_offset, %prefetch_offset_x_t0 : index
%prefetch_offset_x = arith.addi %wg_x_offset, %prefetch_offset_x_t0 : index
%prefetch_offset_y = arith.muli %sg_layout_y, %c32 : index

%k_prefetch_tile = xegpu.create_nd_tdesc %K , shape: [%size_x, %BLOCK_DMODEL], strides: [%BLOCK_DMODEL, %c1] : memref<?x?xf16> -> !xegpu.tensor_desc<16x16xf16, #xegpu.block_tdesc_attr<array_length = 2>>
Expand Down Expand Up @@ -154,15 +156,15 @@ module @flash_attention attributes {gpu.container_module} {
vector<8x16xf32>, vector<8x16xf32>, vector<8x16xf32>, vector<8x16xf32>, vector<8x16xf32>, vector<8x16xf32>, vector<8x16xf32>, vector<8x16xf32>,
vector<8x1xf32>, vector<8x1xf32>, vector<8x1xf32>, vector<8x1xf32>
) {
gpu.barrier

// K prefetch
// K prefetch. Cache hints are required, otherwise the prefetch uses the
// default L1 policy and does not populate L1 for the loads below.
%prefetch_offset_x_running_t = arith.addi %BLOCK_N_3, %k : index
%prefetch_offset_x_running = arith.addi %prefetch_offset_x, %prefetch_offset_x_running_t : index
xegpu.prefetch_nd %k_prefetch_tile[%prefetch_offset_x_running, %prefetch_offset_y] : !xegpu.tensor_desc<16x16xf16, #xegpu.block_tdesc_attr<array_length = 2>>
xegpu.prefetch_nd %k_prefetch_tile[%prefetch_offset_x_running, %prefetch_offset_y] {l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<cached>, l3_hint = #xegpu.cache_hint<cached>} : !xegpu.tensor_desc<16x16xf16, #xegpu.block_tdesc_attr<array_length = 2>>

// V prefetch
xegpu.prefetch_nd %v_prefetch_tile[%prefetch_offset_x_running, %prefetch_offset_y] : !xegpu.tensor_desc<16x16xf16, #xegpu.block_tdesc_attr<array_length = 2>>
xegpu.prefetch_nd %v_prefetch_tile[%prefetch_offset_x_running, %prefetch_offset_y] {l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<cached>, l3_hint = #xegpu.cache_hint<cached>} : !xegpu.tensor_desc<16x16xf16, #xegpu.block_tdesc_attr<array_length = 2>>

// Load first 16x64xf16 (i.e. 16x32xf32) K slice.
%wg_x_offset_running = arith.addi %wg_x_offset, %k : index
Expand Down Expand Up @@ -693,8 +695,8 @@ module @flash_attention attributes {gpu.container_module} {
%magic = arith.constant 0.625 : f32
%c0_f16 = arith.constant 0.0 : f16
%c1_f32 = arith.constant 0.5 : f32
%Z = arith.constant 2 : index // number of batches
%H = arith.constant 2 : index // number of heads
%Z = arith.constant 4 : index // number of batches
%H = arith.constant 4 : index // number of heads
%N_CTX = arith.constant 4096 : index // sequence len
%D_HEAD = arith.constant 64 : index // head dim
%sm_scale = arith.constant 0.5 : f32 // softmax scale
Expand Down
37 changes: 20 additions & 17 deletions test/Integration/Dialect/XeGPU/WG/flash_attention_fwd.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
// RUN: imex-opt %s --gpu-lower-to-xevm-pipeline="xegpu-op-level=workgroup" \
// RUN: imex-opt %s --gpu-lower-to-xevm-pipeline="xegpu-op-level=workgroup igc-cmd-options=-ze-opt-large-register-file" \
// RUN: | mlir-runner \
// RUN: --shared-libs=%mlir_levelzero_runtime \
// RUN: --shared-libs=%mlir_runner_utils \
Expand All @@ -7,6 +7,10 @@
// RUN: --entry-point-result=void \
// RUN: | FileCheck %s

// -ze-opt-large-register-file gives the kernel 256 GRF. At the default 128 GRF
// the loop-carried accumulators do not fit and IGC spills (~8.9 KB, ~282
// spill/fill accesses in the loop body).

#q = #xegpu.layout<sg_layout = [8, 1], sg_data = [16, 64], inst_data = [8, 16]>
#k = #xegpu.layout<sg_layout = [8, 1], sg_data = [16, 64], inst_data = [16, 16]>
#v = #k
Expand Down Expand Up @@ -73,19 +77,20 @@ module @flash_attention attributes {gpu.container_module} {

// K prefetch. Each WG must prefetch 64x64xf16 tile of K per iteration of inner loop.
// For prefetch SG layout is 4x2. Each SG prefetch 16x32xf16 tile.
// Note that prefetch x offset is same as Q x offset. This is because WGs in same batch colloborate on K and V prefetch.
// x offset for prefetch is the same as for the K, V loads (%wg_x_offset),
// so the prefetch tracks the loads BLOCK_N_3 rows ahead.
%k_prefetch_tile = xegpu.create_nd_tdesc %K , shape: [%size_x, %BLOCK_DMODEL], strides: [%BLOCK_DMODEL, %c1] : memref<?x?xf16> -> !xegpu.tensor_desc<64x64xf16, #k_prefetch>
xegpu.prefetch_nd %k_prefetch_tile[%wg_q_x_offset, %c0] {l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<cached>, l3_hint = #xegpu.cache_hint<cached>, layout = #k_prefetch} : !xegpu.tensor_desc<64x64xf16, #k_prefetch>
%wg_q_x_offset_plus_BLOCK_N = arith.addi %wg_q_x_offset, %BLOCK_N : index
xegpu.prefetch_nd %k_prefetch_tile[%wg_q_x_offset_plus_BLOCK_N, %c0] {l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<cached>, l3_hint = #xegpu.cache_hint<cached>, layout = #k_prefetch} : !xegpu.tensor_desc<64x64xf16, #k_prefetch>
%wg_q_x_offset_plus_2_BLOCK_N = arith.addi %wg_q_x_offset_plus_BLOCK_N, %BLOCK_N : index
xegpu.prefetch_nd %k_prefetch_tile[%wg_q_x_offset_plus_2_BLOCK_N, %c0] {l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<cached>, l3_hint = #xegpu.cache_hint<cached>, layout = #k_prefetch} : !xegpu.tensor_desc<64x64xf16, #k_prefetch>
xegpu.prefetch_nd %k_prefetch_tile[%wg_x_offset, %c0] {l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<cached>, l3_hint = #xegpu.cache_hint<cached>, layout = #k_prefetch} : !xegpu.tensor_desc<64x64xf16, #k_prefetch>
%prefetch_offset_x_plus_BLOCK_N = arith.addi %wg_x_offset, %BLOCK_N : index
xegpu.prefetch_nd %k_prefetch_tile[%prefetch_offset_x_plus_BLOCK_N, %c0] {l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<cached>, l3_hint = #xegpu.cache_hint<cached>, layout = #k_prefetch} : !xegpu.tensor_desc<64x64xf16, #k_prefetch>
%prefetch_offset_x_plus_2_BLOCK_N = arith.addi %prefetch_offset_x_plus_BLOCK_N, %BLOCK_N : index
xegpu.prefetch_nd %k_prefetch_tile[%prefetch_offset_x_plus_2_BLOCK_N, %c0] {l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<cached>, l3_hint = #xegpu.cache_hint<cached>, layout = #k_prefetch} : !xegpu.tensor_desc<64x64xf16, #k_prefetch>

// V prefetch is similar to K
%v_prefetch_tile = xegpu.create_nd_tdesc %V , shape: [%size_x, %BLOCK_DMODEL], strides: [%BLOCK_DMODEL, %c1] : memref<?x?xf16> -> !xegpu.tensor_desc<64x64xf16, #v_prefetch>
xegpu.prefetch_nd %v_prefetch_tile[%wg_q_x_offset, %c0] {l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<cached>, l3_hint = #xegpu.cache_hint<cached>, layout = #v_prefetch} : !xegpu.tensor_desc<64x64xf16, #v_prefetch>
xegpu.prefetch_nd %v_prefetch_tile[%wg_q_x_offset_plus_BLOCK_N, %c0] {l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<cached>, l3_hint = #xegpu.cache_hint<cached>, layout = #v_prefetch} : !xegpu.tensor_desc<64x64xf16, #v_prefetch>
xegpu.prefetch_nd %v_prefetch_tile[%wg_q_x_offset_plus_2_BLOCK_N, %c0] {l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<cached>, l3_hint = #xegpu.cache_hint<cached>, layout = #v_prefetch} : !xegpu.tensor_desc<64x64xf16, #v_prefetch>
xegpu.prefetch_nd %v_prefetch_tile[%wg_x_offset, %c0] {l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<cached>, l3_hint = #xegpu.cache_hint<cached>, layout = #v_prefetch} : !xegpu.tensor_desc<64x64xf16, #v_prefetch>
xegpu.prefetch_nd %v_prefetch_tile[%prefetch_offset_x_plus_BLOCK_N, %c0] {l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<cached>, l3_hint = #xegpu.cache_hint<cached>, layout = #v_prefetch} : !xegpu.tensor_desc<64x64xf16, #v_prefetch>
xegpu.prefetch_nd %v_prefetch_tile[%prefetch_offset_x_plus_2_BLOCK_N, %c0] {l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<cached>, l3_hint = #xegpu.cache_hint<cached>, layout = #v_prefetch} : !xegpu.tensor_desc<64x64xf16, #v_prefetch>
%BLOCK_N_3_t = arith.addi %BLOCK_N, %BLOCK_N : index
%BLOCK_N_3 = arith.addi %BLOCK_N_3_t, %BLOCK_N : index

Expand Down Expand Up @@ -119,15 +124,13 @@ module @flash_attention attributes {gpu.container_module} {
-> (
vector<128x64xf32>, vector<128xf32>, vector<128xf32>
) {
gpu.barrier

// K prefetch
%prefetch_offset_x_running_t = arith.addi %BLOCK_N_3, %k : index
%prefetch_offset_x_running = arith.addi %wg_q_x_offset, %prefetch_offset_x_running_t : index
xegpu.prefetch_nd %k_prefetch_tile[%prefetch_offset_x_running, %c0] {layout = #k_prefetch}: !xegpu.tensor_desc<64x64xf16, #k_prefetch>
%prefetch_offset_x_running = arith.addi %wg_x_offset, %prefetch_offset_x_running_t : index
xegpu.prefetch_nd %k_prefetch_tile[%prefetch_offset_x_running, %c0] {layout = #k_prefetch, l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<cached>, l3_hint = #xegpu.cache_hint<cached>}: !xegpu.tensor_desc<64x64xf16, #k_prefetch>

// V prefetch
xegpu.prefetch_nd %v_prefetch_tile[%prefetch_offset_x_running, %c0] {layout = #v_prefetch}: !xegpu.tensor_desc<64x64xf16, #v_prefetch>
xegpu.prefetch_nd %v_prefetch_tile[%prefetch_offset_x_running, %c0] {layout = #v_prefetch, l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<cached>, l3_hint = #xegpu.cache_hint<cached>}: !xegpu.tensor_desc<64x64xf16, #v_prefetch>

// Load first 16x64xf16 K slice. K is in column major layout, so we need to transpose after loading.
%wg_x_offset_running = arith.addi %wg_x_offset, %k : index
Expand Down Expand Up @@ -449,8 +452,8 @@ module @flash_attention attributes {gpu.container_module} {
%magic = arith.constant 0.625 : f32
%c0_f16 = arith.constant 0.0 : f16
%c1_f32 = arith.constant 0.5 : f32
%Z = arith.constant 2 : index // number of batches
%H = arith.constant 2 : index // number of heads
%Z = arith.constant 4 : index // number of batches
%H = arith.constant 4 : index // number of heads
%N_CTX = arith.constant 4096 : index // sequence len
%D_HEAD = arith.constant 64 : index // head dim
%sm_scale = arith.constant 0.5 : f32 // softmax scale
Expand Down
Loading