From 761d121cc5db441dd569e436219fcc0e547d11f0 Mon Sep 17 00:00:00 2001 From: Charitha Saumya Date: Tue, 11 Aug 2026 17:54:22 +0000 Subject: [PATCH 1/3] insert prefetch --- .../xegpu/fused_attention_schedule.py | 60 ++++++++++++++++++- 1 file changed, 58 insertions(+), 2 deletions(-) diff --git a/lighthouse/schedule/xegpu/fused_attention_schedule.py b/lighthouse/schedule/xegpu/fused_attention_schedule.py index 65a2c91e..0d6dfc0e 100644 --- a/lighthouse/schedule/xegpu/fused_attention_schedule.py +++ b/lighthouse/schedule/xegpu/fused_attention_schedule.py @@ -283,8 +283,19 @@ def bundle_xegpu_fused_attention_schedule( wg_loop = loop.loop_forall_to_parallel([anytype], wg_loop) func = transform.get_parent_op(anytype, wg_loop) - # Convert scf.parallel to gpu.launch - func = apply_registered_pass(func, "gpu-map-parallel-loops") + # Convert scf.parallel to gpu.launch. + # + # The parallel loop nest is (batch, rows-within-batch). Map the innermost + # (rows) loop to block X so that the workgroups sharing a batch's K/V are + # consecutive in dispatch order and therefore co-resident. With the default + # outermost-first policy, batch maps to X and the sharers of a batch are + # strided apart by the batch count, which spreads them across dispatch waves + # and defeats the K/V prefetch sharing. + func = apply_registered_pass( + func, + "gpu-map-parallel-loops", + options={"mapping-policy": "innermost-first"}, + ) func = apply_registered_pass(func, "convert-parallel-loops-to-gpu") func = apply_registered_pass(func, "lower-affine") transform.apply_cse(func) @@ -324,6 +335,51 @@ def bundle_xegpu_fused_attention_schedule( transform.apply_cse(gpu_func) gpu_func = apply_registered_pass(gpu_func, "loop-invariant-code-motion") + # Insert prefetches for the K and V tiles of the reduction loop. + # + # The reduction loop is unrolled over the inner tile: there are 4 K loads and + # 4 V loads of 16x64 each, and together the 4 loads of a given operand cover + # the tile_size x n_head region starting at the loop induction variable. Only + # the first load of each group is matched: prefetch_tile_shape overrides the + # 16x64 load shape with the full tile, so one prefetch covers all 4 loads. + # Note this must run before the wg-level layouts are set below, since the + # prefetch descriptor is cloned from the load's descriptor. + prefetch_tile_shape = [tile_size, parameters["n_head"]] + # Distribute the prefetched tile over all subgroups: sg_layout x sg_data must + # cover prefetch_tile_shape exactly ([4, 2] x [16, 32] = [64, 64]). + # TODO: derive these from the tile shape and the subgroup count instead of + # hardcoding them. + prefetch_sg_layout = [4, 2] + prefetch_sg_data = [16, 32] + prefetch_inst_data = [16, 16] + # 9 load_nd ops: index 0 is the Q load (hoisted out of the loop), 1-4 are the + # K loads and 5-8 are the V loads, in program order. + initial_load_nd_ops = match_and_split(gpu_func, ops={"xegpu.load_nd"}, nhandles=9) + q_load = initial_load_nd_ops[0] + first_k_load = initial_load_nd_ops[1] + first_v_load = initial_load_nd_ops[5] + for load_op in [first_k_load, first_v_load]: + # Base the prefetch at the Q load's offsets. Those are dynamic values + # (derived from the block id), so they can only be supplied by handle via + # offsets_from, not as a static index list. The loop dimension still + # advances by the loop step on top of this base. + desc_op = xegpu.insert_prefetch( + load_op, + nb_prefetch=parameters.get("nb_prefetch", 3), + prefetch_tile_shape=prefetch_tile_shape, + offsets_from=q_load, + ) + # The emitted prefetch_nd ops are the consumers of the new descriptor. + prefetch_ops = transform.get_consumers_of_result(anytype, desc_op, 0) + xegpu.set_anchor_layout( + prefetch_ops, + sg_layout=prefetch_sg_layout, + sg_data=prefetch_sg_data, + inst_data=prefetch_inst_data, + ) + transform.apply_cse(gpu_func) + canonicalize(gpu_func) + if stop_at_stage == "xegpu-initial": raise PipelineInterrupt() From d306ec930b9217dbcfb94e82fb697b22feddc527 Mon Sep 17 00:00:00 2001 From: Charitha Saumya Date: Thu, 13 Aug 2026 05:58:12 +0000 Subject: [PATCH 2/3] simplify fused attention code gen + add fastmath + add prefetch --- examples/xegpu/fused_attention.py | 23 +- examples/xegpu/nanoGPT_schedule.py | 132 ++++---- .../ops/replace_with_fused_attention.py | 317 ++++++++---------- .../xegpu/fused_attention_schedule.py | 274 +++++++-------- 4 files changed, 337 insertions(+), 409 deletions(-) diff --git a/examples/xegpu/fused_attention.py b/examples/xegpu/fused_attention.py index e3608126..7d20cfd8 100644 --- a/examples/xegpu/fused_attention.py +++ b/examples/xegpu/fused_attention.py @@ -27,14 +27,16 @@ def fused_attention_complexity(Z: int, H: int, n_ctx: int, n_head: int, nbytes: """ Complexity of fused attention operation. - For each batch and head: - - Q @ K^T: O(n_ctx^2 * n_head) operations - - Softmax: O(n_ctx^2) operations - - Attention @ V: O(n_ctx^2 * n_head) operations - Total: approximately 2*n_ctx^2*n_head FLOPs per batch and head + Counts the two matmuls only, at 2 FLOPs per multiply-accumulate, which is the + convention used by the flash attention tutorials (and hence what published + attention FLOPS numbers can be compared against). For each batch and head: + - 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 + is not multiply-accumulate work. Halve the total for a causal mask. """ - # Approximation: 2 * n_ctx^2 * n_head FLOPs per batch and head - flop_count = Z * H * 2 * n_ctx * n_ctx * n_head + # 2 matmuls, 2 * n_ctx^2 * n_head FLOPs each, per batch and head + flop_count = Z * H * 4 * n_ctx * n_ctx * n_head # Memory: read Q, K, V and write output memory_reads = 3 * Z * H * n_ctx * n_head * nbytes memory_writes = Z * H * n_ctx * n_head * nbytes @@ -251,6 +253,12 @@ def parse_cli(): default=64, help="Tile size for the inner reduction dimension (K/V sequence length)", ) + parser.add_argument( + "--nb-prefetch", + type=int, + default=1, + help="Number of K/V tiles to prefetch ahead in the inner loop (0 disables).", + ) parser.add_argument( "--nruns", type=int, @@ -313,6 +321,7 @@ def parse_cli(): "sg_rows": args.sg_rows, "subgroup_size": args.subgroup_size, "inner_loop_tile_size": args.inner_loop_tile_size, + "nb_prefetch": args.nb_prefetch, } Z = args.batch_size diff --git a/examples/xegpu/nanoGPT_schedule.py b/examples/xegpu/nanoGPT_schedule.py index 9e15a6aa..539f7d7f 100644 --- a/examples/xegpu/nanoGPT_schedule.py +++ b/examples/xegpu/nanoGPT_schedule.py @@ -254,26 +254,24 @@ def xegpu_fa_annotation(gf, anytype, fa_params): """Attach XeGPU layouts to one fused-attention gpu.func.""" num_subgroups = fa_params["wg_rows"] // fa_params["sg_rows"] n_head = fa_params["n_head"] + tile_size = fa_params["inner_loop_tile_size"] q_sg_layout = [num_subgroups, 1] q_sg_data = [16, n_head] q_inst_data = [8, 16] + # K and V tiles are [tile_size, n_head], shared by all subgroups. k_sg_layout = [num_subgroups, 1] - k_sg_data = [16, n_head] + k_sg_data = [tile_size, n_head] k_inst_data = [16, 16] v_sg_layout, v_sg_data, v_inst_data = k_sg_layout, k_sg_data, k_inst_data kt_sg_layout = [1, num_subgroups] - kt_sg_data = [n_head, 16] + kt_sg_data = [n_head, tile_size] kt_inst_data = [16, 16] kt_order = [0, 1] out_sg_layout, out_sg_data, out_inst_data = q_sg_layout, q_sg_data, q_inst_data - l128_sg_layout = [num_subgroups, 1] - l128_sg_data = [16, 16] - l128_inst_data = [8, 16] - qk_sg_layout, qk_sg_data, qk_inst_data = ( - l128_sg_layout, - l128_sg_data, - l128_inst_data, - ) + # Q@K^T (attention weights) tile is [wg_rows, tile_size]. + qk_sg_layout = [num_subgroups, 1] + qk_sg_data = [16, tile_size] + qk_inst_data = [8, 16] store_nd_op = match_and_split(gf, ops={"xegpu.store_nd"}, nhandles=1)[0] xegpu.set_anchor_layout( @@ -282,64 +280,68 @@ def xegpu_fa_annotation(gf, anytype, fa_params): sg_data=out_sg_data, inst_data=out_inst_data, ) - load_nd_ops = match_and_split(gf, ops={"xegpu.load_nd"}, nhandles=9) + # 3 load_nd ops: Q (hoisted out of the loop), then K and V in the loop. + load_nd_ops = match_and_split(gf, ops={"xegpu.load_nd"}, nhandles=3) xegpu.set_anchor_layout( load_nd_ops[0], sg_layout=q_sg_layout, sg_data=q_sg_data, inst_data=q_inst_data ) - for i in range(1, 5): - xegpu.set_anchor_layout( - load_nd_ops[i], - sg_layout=k_sg_layout, - sg_data=k_sg_data, - inst_data=k_inst_data, - ) - for i in range(5, 9): - xegpu.set_anchor_layout( - load_nd_ops[i], - sg_layout=v_sg_layout, - sg_data=v_sg_data, - inst_data=v_inst_data, - ) - dpas_ops = match_and_split(gf, ops={"xegpu.dpas"}, nhandles=8) - for i in range(4): - d = dpas_ops[i] - xegpu.set_anchor_layout( - d, sg_layout=q_sg_layout, sg_data=q_sg_data, inst_data=q_inst_data, index=0 - ) - xegpu.set_anchor_layout( - d, - sg_layout=kt_sg_layout, - sg_data=kt_sg_data, - inst_data=kt_inst_data, - order=kt_order, - index=1, - ) - xegpu.set_anchor_layout( - d, - sg_layout=l128_sg_layout, - sg_data=l128_sg_data, - inst_data=l128_inst_data, - index=2, - ) - for i in range(4, 8): - d = dpas_ops[i] - xegpu.set_anchor_layout( - d, - sg_layout=qk_sg_layout, - sg_data=qk_sg_data, - inst_data=qk_inst_data, - index=0, - ) - xegpu.set_anchor_layout( - d, sg_layout=v_sg_layout, sg_data=v_sg_data, inst_data=v_inst_data, index=1 - ) - xegpu.set_anchor_layout( - d, - sg_layout=out_sg_layout, - sg_data=out_sg_data, - inst_data=out_inst_data, - index=2, - ) + xegpu.set_anchor_layout( + load_nd_ops[1], + sg_layout=k_sg_layout, + sg_data=k_sg_data, + inst_data=k_inst_data, + ) + xegpu.set_anchor_layout( + load_nd_ops[2], + sg_layout=v_sg_layout, + sg_data=v_sg_data, + inst_data=v_inst_data, + ) + # 2 dpas ops: Q@K^T and P@V. + qk_dpas, pv_dpas = match_and_split(gf, ops={"xegpu.dpas"}, nhandles=2) + xegpu.set_anchor_layout( + qk_dpas, + sg_layout=q_sg_layout, + sg_data=q_sg_data, + inst_data=q_inst_data, + index=0, + ) + xegpu.set_anchor_layout( + qk_dpas, + sg_layout=kt_sg_layout, + sg_data=kt_sg_data, + inst_data=kt_inst_data, + order=kt_order, + index=1, + ) + xegpu.set_anchor_layout( + qk_dpas, + sg_layout=qk_sg_layout, + sg_data=qk_sg_data, + inst_data=qk_inst_data, + index=2, + ) + xegpu.set_anchor_layout( + pv_dpas, + sg_layout=qk_sg_layout, + sg_data=qk_sg_data, + inst_data=qk_inst_data, + index=0, + ) + xegpu.set_anchor_layout( + pv_dpas, + sg_layout=v_sg_layout, + sg_data=v_sg_data, + inst_data=v_inst_data, + index=1, + ) + xegpu.set_anchor_layout( + pv_dpas, + sg_layout=out_sg_layout, + sg_data=out_sg_data, + inst_data=out_inst_data, + index=2, + ) def build_combined_schedule( diff --git a/lighthouse/dialects/transform/transform_ext/ops/replace_with_fused_attention.py b/lighthouse/dialects/transform/transform_ext/ops/replace_with_fused_attention.py index 570545c1..4e22eaf7 100644 --- a/lighthouse/dialects/transform/transform_ext/ops/replace_with_fused_attention.py +++ b/lighthouse/dialects/transform/transform_ext/ops/replace_with_fused_attention.py @@ -17,22 +17,21 @@ def emit_vector_constant(shape, fill_value, element_type): return arith.constant(vector_type, attr) -def compute_qkt_chunks( +def compute_qkt( q_value, k_load_op, loop_idx, - k_tile_offsets, wg_rows, d_head, - k_subtile_size, - num_k_tiles, + tile_size, element_type, + compute_type, ): - """Load K tiles, transpose, and contract with Q to produce Q@K^T chunks. + """Load the K tile, transpose it, and contract with Q to produce Q@K^T. - Each K tile is [k_subtile_size, d_head], transposed to [d_head, k_subtile_size] - and contracted with q_value [wg_rows, d_head] to produce [wg_rows, k_subtile_size]. - Returns a list of `num_k_tiles` such chunks. + The K tile is [tile_size, d_head], transposed to [d_head, tile_size] and + contracted with q_value [wg_rows, d_head] to produce [wg_rows, tile_size]. + Q and K are `element_type`, the contraction accumulates in `compute_type`. """ k_memref = k_load_op.operands[0] k_load_indices = list(k_load_op.operands[1:-1]) @@ -64,141 +63,98 @@ def compute_qkt_chunks( ] ) - qkt_chunk_type = ir.VectorType.get([wg_rows, k_subtile_size], element_type) - qkt_chunk_acc = emit_vector_constant((wg_rows, k_subtile_size), 0.0, element_type) - - qkt_chunks = [] - for tile_idx in range(num_k_tiles): - k_tile_offset = arith.addi(loop_idx, k_tile_offsets[tile_idx]) - - k_tile_indices = k_load_indices.copy() - k_tile_indices[-2] = k_tile_offset - - k_tile_type = ir.VectorType.get([k_subtile_size, d_head], element_type) - k_tile = vector.TransferReadOp( - k_tile_type, - k_memref, - k_tile_indices, - k_perm_map, - padding, - in_bounds=in_bounds, - ).result - - k_transpose_type = ir.VectorType.get([d_head, k_subtile_size], element_type) - k_transpose = vector.transpose(k_transpose_type, k_tile, [1, 0]) - - qkt_chunk = vector.contract( - qkt_chunk_type, - q_value, - k_transpose, - qkt_chunk_acc, - indexing_maps=indexing_maps, - iterator_types=iterator_types, - ) - qkt_chunks.append(qkt_chunk) - - return qkt_chunks - - -def compute_qkt_max_scaled(qkt_chunks, num_k_tiles, m_i_init, scale_vector): - """Reduce Q@K^T chunks to a row-wise scaled max. - - Combines chunks elementwise with maximumf, reduces along the inner dim with - multi_reduction(maxnumf, acc=m_i_init), and multiplies by scale_vector. - Returns a [wg_rows] vector. - """ - qkt_max_combined = qkt_chunks[0] - for i in range(1, num_k_tiles): - qkt_max_combined = arith.maximumf(qkt_max_combined, qkt_chunks[i]) - - qkt_max = vector.multi_reduction( - kind="maxnumf", - source=qkt_max_combined, - acc=m_i_init, - reduction_dims=[1], + qkt_type = ir.VectorType.get([wg_rows, tile_size], compute_type) + qkt_acc = emit_vector_constant((wg_rows, tile_size), 0.0, compute_type) + + k_tile_indices = k_load_indices.copy() + k_tile_indices[-2] = loop_idx + + k_tile_type = ir.VectorType.get([tile_size, d_head], element_type) + k_tile = vector.TransferReadOp( + k_tile_type, + k_memref, + k_tile_indices, + k_perm_map, + padding, + in_bounds=in_bounds, + ).result + + k_transpose_type = ir.VectorType.get([d_head, tile_size], element_type) + k_transpose = vector.transpose(k_transpose_type, k_tile, [1, 0]) + + return vector.contract( + qkt_type, + q_value, + k_transpose, + qkt_acc, + indexing_maps=indexing_maps, + iterator_types=iterator_types, ) - return arith.mulf(qkt_max, scale_vector) - def compute_online_softmax_and_sum( - qkt_chunks, + qkt_scaled, m_ij, l_i_init, - scale_value, wg_rows, - k_subtile_size, - num_k_tiles, - element_type, + tile_size, + compute_type, ): - """Apply online softmax to Q@K^T chunks and reduce to a row-wise sum. + """Apply online softmax to the scaled Q@K^T and reduce to a row-wise sum. - For each chunk: exp(chunk * scale - m_ij), broadcast over the inner dim. - Returns (qkt_exp_chunks, l_ij) where qkt_exp_chunks is the list of - [wg_rows, k_subtile_size] exp tiles and l_ij is their row-wise sum - [wg_rows] (added into l_i_init). + Computes exp(qkt_scaled - m_ij), with m_ij broadcast over the inner dim. + Returns (qkt_exp, l_ij) where qkt_exp is the [wg_rows, tile_size] exp tile + and l_ij is its row-wise sum [wg_rows] (added into l_i_init). """ - scale_chunk = emit_vector_constant( - (wg_rows, k_subtile_size), scale_value, element_type - ) - - # Broadcast m_ij from [wg_rows] to [wg_rows, k_subtile_size] - m_ij_bcasted_type = ir.VectorType.get([k_subtile_size, wg_rows], element_type) + # Broadcast m_ij from [wg_rows] to [wg_rows, tile_size] + m_ij_bcasted_type = ir.VectorType.get([tile_size, wg_rows], compute_type) m_ij_bcasted = vector.broadcast(m_ij_bcasted_type, m_ij) - m_ij_transposed_type = ir.VectorType.get([wg_rows, k_subtile_size], element_type) + m_ij_transposed_type = ir.VectorType.get([wg_rows, tile_size], compute_type) m_ij_transposed = vector.transpose(m_ij_transposed_type, m_ij_bcasted, [1, 0]) - qkt_exp_chunks = [] - for qkt_chunk in qkt_chunks: - qkt_scaled = arith.mulf(qkt_chunk, scale_chunk) - qkt_centered = arith.subf(qkt_scaled, m_ij_transposed) - qkt_exp = math.exp(qkt_centered) - qkt_exp_chunks.append(qkt_exp) - - qkt_exp_combined = qkt_exp_chunks[0] - for i in range(1, num_k_tiles): - qkt_exp_combined = arith.addf(qkt_exp_combined, qkt_exp_chunks[i]) + qkt_centered = arith.subf(qkt_scaled, m_ij_transposed) + # fastmath lets the exp lower to the native hardware exp; without it + # the accurate expansion doubles the exp count and scalarizes part of it. + qkt_exp = math.exp(qkt_centered, fastmath="fast") l_ij = vector.multi_reduction( kind="add", - source=qkt_exp_combined, + source=qkt_exp, acc=l_i_init, reduction_dims=[1], ) - return qkt_exp_chunks, l_ij + return qkt_exp, l_ij -def rescale_pv_out_accumulator(acc, alpha, wg_rows, d_head, element_type): +def rescale_pv_out_accumulator(acc, alpha, wg_rows, d_head, compute_type): """Rescale the running P@V accumulator by broadcasting alpha across d_head. Broadcasts alpha [wg_rows] to [wg_rows, d_head] and multiplies acc by it elementwise. Returns the rescaled accumulator. """ - alpha_bcasted_type = ir.VectorType.get([d_head, wg_rows], element_type) + alpha_bcasted_type = ir.VectorType.get([d_head, wg_rows], compute_type) alpha_bcasted = vector.broadcast(alpha_bcasted_type, alpha) - alpha_transposed_type = ir.VectorType.get([wg_rows, d_head], element_type) + alpha_transposed_type = ir.VectorType.get([wg_rows, d_head], compute_type) alpha_transposed = vector.transpose(alpha_transposed_type, alpha_bcasted, [1, 0]) return arith.mulf(acc, alpha_transposed) -def compute_pv_chunks( - qkt_exp_chunks, +def compute_pv( + qkt_exp, v_load_op, pv_init, loop_idx, - k_tile_offsets, acc_vector_type, d_head, - k_subtile_size, - num_k_tiles, + tile_size, element_type, ): - """Load V tiles and contract with softmax chunks, accumulating into pv_init. + """Load the V tile and contract it with the softmax tile, accumulating into pv_init. - For each tile: load V tile [k_subtile_size, d_head] and contract with the - matching exp chunk [wg_rows, k_subtile_size] into the running [wg_rows, d_head] - accumulator. Returns the final accumulated result. + Loads V [tile_size, d_head] (`element_type`) and contracts it with the exp + tile [wg_rows, tile_size] into the running [wg_rows, d_head] accumulator, + whose type is given by `acc_vector_type`. Returns the accumulated result. """ v_memref = v_load_op.operands[0] v_load_indices = list(v_load_op.operands[1:-1]) @@ -229,40 +185,34 @@ def compute_pv_chunks( ] ) - pv_out = pv_init - for tile_idx in range(num_k_tiles): - v_tile_offset = arith.addi(loop_idx, k_tile_offsets[tile_idx]) - - v_tile_indices = v_load_indices.copy() - v_tile_indices[-2] = v_tile_offset - - v_tile_type = ir.VectorType.get([k_subtile_size, d_head], element_type) - v_tile = vector.TransferReadOp( - v_tile_type, - v_memref, - v_tile_indices, - v_perm_map, - v_padding, - in_bounds=v_in_bounds, - ).result - - pv_out = vector.contract( - acc_vector_type, - qkt_exp_chunks[tile_idx], - v_tile, - pv_out, - indexing_maps=indexing_maps_pv, - iterator_types=iterator_types_pv, - ) - - return pv_out + v_tile_indices = v_load_indices.copy() + v_tile_indices[-2] = loop_idx + + v_tile_type = ir.VectorType.get([tile_size, d_head], element_type) + v_tile = vector.TransferReadOp( + v_tile_type, + v_memref, + v_tile_indices, + v_perm_map, + v_padding, + in_bounds=v_in_bounds, + ).result + + return vector.contract( + acc_vector_type, + qkt_exp, + v_tile, + pv_init, + indexing_maps=indexing_maps_pv, + iterator_types=iterator_types_pv, + ) -def normalize_ouput_by_sum(pv_out, l_i_out, wg_rows, d_head, element_type): +def normalize_ouput_by_sum(pv_out, l_i_out, wg_rows, d_head, compute_type): """Divide pv_out [wg_rows, d_head] by l_i_out [wg_rows] (broadcast over d_head).""" - l_i_out_bcasted_type = ir.VectorType.get([d_head, wg_rows], element_type) + l_i_out_bcasted_type = ir.VectorType.get([d_head, wg_rows], compute_type) l_i_out_bcasted = vector.broadcast(l_i_out_bcasted_type, l_i_out) - l_i_out_transposed_type = ir.VectorType.get([wg_rows, d_head], element_type) + l_i_out_transposed_type = ir.VectorType.get([wg_rows, d_head], compute_type) l_i_out_transposed = vector.transpose( l_i_out_transposed_type, l_i_out_bcasted, [1, 0] ) @@ -374,34 +324,32 @@ def apply( # Get tile size tile_size_value = ir.IntegerAttr(op.tile_size).value - # Get element type from q_load result + # Get element type from q_load result. Q, K, V and the output stay in this + # type (the DPAS input type), while the softmax and both accumulators are + # computed in f32 to keep the online softmax accurate and to avoid holding + # the intermediates in narrow registers. element_type = q_vector_type.element_type + compute_type = ir.F32Type.get() # Build the fused attention computation with ir.InsertionPoint(output_op): # Define m_i_init: vector of shape [wg_rows] with neg_inf values - # NOTE: We use float32 for the initial neg_inf values and cast to the element type - # to avoid issues with representing -inf. - m_i_vector_type = ir.VectorType.get([wg_rows], element_type) - m_i_init_f32 = emit_vector_constant( - (wg_rows,), float("-inf"), ir.F32Type.get() - ) - m_i_init = arith.truncf(m_i_vector_type, m_i_init_f32) + m_i_init = emit_vector_constant((wg_rows,), float("-inf"), compute_type) # Define l_i_init: vector of shape [wg_rows] with zero values - l_i_init = emit_vector_constant((wg_rows,), 0.0, element_type) + l_i_init = emit_vector_constant((wg_rows,), 0.0, compute_type) # Define acc_init: vector of shape [wg_rows, d_head] with zero values - acc_vector_type = ir.VectorType.get([wg_rows, d_head], element_type) - acc_init = emit_vector_constant((wg_rows, d_head), 0.0, element_type) + acc_vector_type = ir.VectorType.get([wg_rows, d_head], compute_type) + acc_init = emit_vector_constant((wg_rows, d_head), 0.0, compute_type) # Get n_ctx from k_load result type (first dimension size) k_load_result = k_load_op.results[0] k_vector_type = ir.VectorType(k_load_result.type) n_ctx = k_vector_type.shape[0] - # Define scale vector: vector of shape [wg_rows] with the scale value - scale_vector = emit_vector_constant( - (wg_rows,), scale_value, element_type + # Define scale tile: [wg_rows, tile_size] filled with the scale value + scale_tile = emit_vector_constant( + (wg_rows, tile_size_value), scale_value, compute_type ) # Create loop bounds @@ -424,53 +372,44 @@ def apply( q_value = q_load_op.results[0] - # Constants for K/V tiling (tile into chunks of 16) - k_subtile_size = 16 - num_k_tiles = tile_size_value // k_subtile_size - - # Create offset constants for each K tile - k_tile_offsets = [] - for i in range(num_k_tiles): - offset = arith.constant(index_type, i * k_subtile_size) - k_tile_offsets.append(offset) - - # Load K tiles, transpose, and contract with Q to get Q@K^T chunks - qkt_chunks = compute_qkt_chunks( + # Load the K tile, transpose it, and contract with Q to get Q@K^T + qkt = compute_qkt( q_value, k_load_op, loop_idx, - k_tile_offsets, wg_rows, d_head, - k_subtile_size, - num_k_tiles, + tile_size_value, element_type, + compute_type, ) - - # Reduce Q@K^T chunks to row-wise scaled max: [wg_rows] - qkt_max_scaled = compute_qkt_max_scaled( - qkt_chunks, num_k_tiles, m_i_init, scale_vector + qkt_scaled = arith.mulf(qkt, scale_tile) + + # Reduce the scaled Q@K^T to a row-wise max: [wg_rows] + qkt_row_max = vector.multi_reduction( + kind="maximumf", + source=qkt_scaled, + acc=m_i_init, + reduction_dims=[1], ) - # Compute m_ij = max(m_i, qkt_max_scaled) + # Compute m_ij = max(m_i, qkt_row_max) # Both have shape [wg_rows] - m_ij = arith.maximumf(m_i, qkt_max_scaled) + m_ij = arith.maximumf(m_i, qkt_row_max) - # Apply online softmax to chunks and reduce to row-wise sum - qkt_exp_chunks, l_ij = compute_online_softmax_and_sum( - qkt_chunks, + # Apply online softmax and reduce to row-wise sum + qkt_exp, l_ij = compute_online_softmax_and_sum( + qkt_scaled, m_ij, l_i_init, - scale_value, wg_rows, - k_subtile_size, - num_k_tiles, - element_type, + tile_size_value, + compute_type, ) # Compute alpha = exp(m_i - m_ij) m_diff = arith.subf(m_i, m_ij) - alpha = math.exp(m_diff) + alpha = math.exp(m_diff, fastmath="fast") # Update l_i: l_i_updated = l_i * alpha + l_ij l_i_scaled = arith.mulf(l_i, alpha) @@ -478,20 +417,24 @@ def apply( # Rescale running P@V accumulator by alpha acc_updated = rescale_pv_out_accumulator( - acc, alpha, wg_rows, d_head, element_type + acc, alpha, wg_rows, d_head, compute_type ) - # Load V tiles and contract with softmax chunks into pv_out - pv_out = compute_pv_chunks( - qkt_exp_chunks, + # Narrow the softmax tile to the DPAS input type + qkt_exp_type = ir.VectorType.get( + [wg_rows, tile_size_value], element_type + ) + qkt_exp_narrow = arith.truncf(qkt_exp_type, qkt_exp) + + # Load the V tile and contract with the softmax tile into pv_out + pv_out = compute_pv( + qkt_exp_narrow, v_load_op, acc_updated, loop_idx, - k_tile_offsets, acc_vector_type, d_head, - k_subtile_size, - num_k_tiles, + tile_size_value, element_type, ) @@ -503,8 +446,12 @@ def apply( l_i_out = loop.results[1] with ir.InsertionPoint.after(loop): # Normalize the output: output_final = pv_out / l_i_out - output_final = normalize_ouput_by_sum( - pv_out, l_i_out, wg_rows, d_head, element_type + output_normalized = normalize_ouput_by_sum( + pv_out, l_i_out, wg_rows, d_head, compute_type + ) + # Narrow back to the type of the output op being replaced + output_final = arith.truncf( + output_op.results[0].type, output_normalized ) # Replace all uses of the original output operation with the final loop result diff --git a/lighthouse/schedule/xegpu/fused_attention_schedule.py b/lighthouse/schedule/xegpu/fused_attention_schedule.py index 0d6dfc0e..f79442bf 100644 --- a/lighthouse/schedule/xegpu/fused_attention_schedule.py +++ b/lighthouse/schedule/xegpu/fused_attention_schedule.py @@ -283,19 +283,8 @@ def bundle_xegpu_fused_attention_schedule( wg_loop = loop.loop_forall_to_parallel([anytype], wg_loop) func = transform.get_parent_op(anytype, wg_loop) - # Convert scf.parallel to gpu.launch. - # - # The parallel loop nest is (batch, rows-within-batch). Map the innermost - # (rows) loop to block X so that the workgroups sharing a batch's K/V are - # consecutive in dispatch order and therefore co-resident. With the default - # outermost-first policy, batch maps to X and the sharers of a batch are - # strided apart by the batch count, which spreads them across dispatch waves - # and defeats the K/V prefetch sharing. - func = apply_registered_pass( - func, - "gpu-map-parallel-loops", - options={"mapping-policy": "innermost-first"}, - ) + # Convert scf.parallel to gpu.launch + func = apply_registered_pass(func, "gpu-map-parallel-loops") func = apply_registered_pass(func, "convert-parallel-loops-to-gpu") func = apply_registered_pass(func, "lower-affine") transform.apply_cse(func) @@ -335,84 +324,62 @@ def bundle_xegpu_fused_attention_schedule( transform.apply_cse(gpu_func) gpu_func = apply_registered_pass(gpu_func, "loop-invariant-code-motion") - # Insert prefetches for the K and V tiles of the reduction loop. - # - # The reduction loop is unrolled over the inner tile: there are 4 K loads and - # 4 V loads of 16x64 each, and together the 4 loads of a given operand cover - # the tile_size x n_head region starting at the loop induction variable. Only - # the first load of each group is matched: prefetch_tile_shape overrides the - # 16x64 load shape with the full tile, so one prefetch covers all 4 loads. - # Note this must run before the wg-level layouts are set below, since the - # prefetch descriptor is cloned from the load's descriptor. - prefetch_tile_shape = [tile_size, parameters["n_head"]] - # Distribute the prefetched tile over all subgroups: sg_layout x sg_data must - # cover prefetch_tile_shape exactly ([4, 2] x [16, 32] = [64, 64]). - # TODO: derive these from the tile shape and the subgroup count instead of - # hardcoding them. - prefetch_sg_layout = [4, 2] - prefetch_sg_data = [16, 32] - prefetch_inst_data = [16, 16] - # 9 load_nd ops: index 0 is the Q load (hoisted out of the loop), 1-4 are the - # K loads and 5-8 are the V loads, in program order. - initial_load_nd_ops = match_and_split(gpu_func, ops={"xegpu.load_nd"}, nhandles=9) - q_load = initial_load_nd_ops[0] - first_k_load = initial_load_nd_ops[1] - first_v_load = initial_load_nd_ops[5] - for load_op in [first_k_load, first_v_load]: - # Base the prefetch at the Q load's offsets. Those are dynamic values - # (derived from the block id), so they can only be supplied by handle via - # offsets_from, not as a static index list. The loop dimension still - # advances by the loop step on top of this base. - desc_op = xegpu.insert_prefetch( - load_op, - nb_prefetch=parameters.get("nb_prefetch", 3), - prefetch_tile_shape=prefetch_tile_shape, - offsets_from=q_load, - ) - # The emitted prefetch_nd ops are the consumers of the new descriptor. - prefetch_ops = transform.get_consumers_of_result(anytype, desc_op, 0) - xegpu.set_anchor_layout( - prefetch_ops, - sg_layout=prefetch_sg_layout, - sg_data=prefetch_sg_data, - inst_data=prefetch_inst_data, + # Insert prefetches for the K and V tiles of the reduction loop. Each inserts + # nb_prefetch prefetches ahead of the loop plus one per iteration, at + # induction_var + nb_prefetch * step. This must run before the wg-level layouts + # are set below, since the prefetch descriptor is cloned from the load's + # descriptor. The layouts of the emitted prefetch_nd ops are set together with + # the other wg-level layouts. + nb_prefetch = parameters.get("nb_prefetch", 1) + if nb_prefetch > 0: + # 3 load_nd ops: Q (hoisted out of the loop), then K and V inside it. + initial_load_nd_ops = match_and_split( + gpu_func, ops={"xegpu.load_nd"}, nhandles=3 ) - transform.apply_cse(gpu_func) - canonicalize(gpu_func) + for load_op in initial_load_nd_ops[1:]: + xegpu.insert_prefetch(load_op, nb_prefetch=nb_prefetch) + transform.apply_cse(gpu_func) + canonicalize(gpu_func) if stop_at_stage == "xegpu-initial": raise PipelineInterrupt() # Define XeGPU layout parameters n_head = parameters["n_head"] - q_sg_layout = [num_subgroups, 1] - q_sg_data = [16, n_head] - q_inst_data = [8, 16] - - k_sg_layout = [num_subgroups, 1] - k_sg_data = [16, n_head] - k_inst_data = [16, 16] - - v_sg_layout = k_sg_layout - v_sg_data = k_sg_data - v_inst_data = k_inst_data + sg_rows = parameters["sg_rows"] - kt_sg_layout = [1, num_subgroups] - kt_sg_data = [n_head, 16] - kt_inst_data = [16, 16] - kt_order = [0, 1] + # Q, the attention weights and the output are all [wg_rows, n_head] tiles that + # are split by rows over the subgroups. Only the memory ops carry inst_data, + # the DPAS operands are left to the default DPAS blocking. + q_sg_layout = [num_subgroups, 1] + q_sg_data = [sg_rows, n_head] + q_load_inst_data = [16, 32] out_sg_layout = q_sg_layout out_sg_data = q_sg_data - out_inst_data = q_inst_data - - layout_128x16_sg_layout = [num_subgroups, 1] - layout_128x16_sg_data = [16, 16] - layout_128x16_inst_data = [8, 16] - qk_sg_layout = layout_128x16_sg_layout - qk_sg_data = layout_128x16_sg_data - qk_inst_data = layout_128x16_inst_data + qk_sg_layout = q_sg_layout + qk_sg_data = [sg_rows, tile_size] + + # The K and V tiles are consumed in full by every subgroup (each subgroup owns + # its own rows of Q). + kv_sg_layout = [1, 1] + kv_load_sg_data = [tile_size, n_head] + v_load_inst_data = [32, 32] + # Load K column-major so that the transpose feeding the DPAS is a no-op. + k_load_order = [0, 1] + + # K^T operand of the Q@K^T DPAS: [n_head, tile_size] + kt_sg_data = [n_head, tile_size] + # V operand of the P@V DPAS: [tile_size, n_head] + v_sg_data = [tile_size, n_head] + + # The K/V prefetches cover the same [tile_size, n_head] tile as the loads, but + # are distributed over the subgroups. [4, 2] layout is used to maximize + # prefetch bandwidth. + prefetch_sg_layout = [4, 2] + prefetch_sg_data = [tile_size // 4, n_head // 2] + prefetch_inst_data = list(prefetch_sg_data) # Set layout attributes for xegpu.store_nd ops. store_nd_op = match_and_split(gpu_func, ops={"xegpu.store_nd"}, nhandles=1)[0] @@ -420,92 +387,95 @@ def bundle_xegpu_fused_attention_schedule( store_nd_op, sg_layout=out_sg_layout, sg_data=out_sg_data, - inst_data=out_inst_data, ) - # Set layout for xegpu.load_nd ops (9 total: 1 Q, 4 K, 4 V) - load_nd_ops = match_and_split(gpu_func, ops={"xegpu.load_nd"}, nhandles=9) + # Set layout for xegpu.load_nd ops (3 total: Q, K, V) + load_nd_ops = match_and_split(gpu_func, ops={"xegpu.load_nd"}, nhandles=3) # First load_nd: Q layout xegpu.set_anchor_layout( - load_nd_ops[0], sg_layout=q_sg_layout, sg_data=q_sg_data, inst_data=q_inst_data + load_nd_ops[0], + sg_layout=q_sg_layout, + sg_data=q_sg_data, + inst_data=q_load_inst_data, ) - # Next 4 load_nd ops: K layout - for load_op in load_nd_ops[:4]: - xegpu.set_anchor_layout( - load_op, - sg_layout=k_sg_layout, - sg_data=k_sg_data, - inst_data=k_inst_data, - ) + # Second load_nd: K layout + xegpu.set_anchor_layout( + load_nd_ops[1], + sg_layout=kv_sg_layout, + sg_data=kv_load_sg_data, + order=k_load_order, + ) - # Last 4 load_nd ops: V layout - for load_op in load_nd_ops[4:]: + # Third load_nd: V layout + xegpu.set_anchor_layout( + load_nd_ops[2], + sg_layout=kv_sg_layout, + sg_data=kv_load_sg_data, + inst_data=v_load_inst_data, + ) + + # Set layout for all K/V xegpu.prefetch_nd ops + if nb_prefetch > 0: + prefetch_ops = match(gpu_func, ops={"xegpu.prefetch_nd"}) xegpu.set_anchor_layout( - load_op, - sg_layout=v_sg_layout, - sg_data=v_sg_data, - inst_data=v_inst_data, + prefetch_ops, + sg_layout=prefetch_sg_layout, + sg_data=prefetch_sg_data, + inst_data=prefetch_inst_data, ) - # Set layout for xegpu.dpas ops (8 total: 4 for Q@K, 4 for P@V) - dpas_ops = match_and_split(gpu_func, ops={"xegpu.dpas"}, nhandles=8) + # Set layout for xegpu.dpas ops (2 total: Q@K^T and P@V) + dpas_ops = match_and_split(gpu_func, ops={"xegpu.dpas"}, nhandles=2) - # Layouts for first 4 dpas ops (Q@K^T): - for qk_dpas_op in dpas_ops[:4]: - # Index 0: Q layout - xegpu.set_anchor_layout( - qk_dpas_op, - sg_layout=q_sg_layout, - sg_data=q_sg_data, - inst_data=q_inst_data, - index=0, - ) - # Index 1: K^T layout - xegpu.set_anchor_layout( - qk_dpas_op, - sg_layout=kt_sg_layout, - sg_data=kt_sg_data, - inst_data=kt_inst_data, - order=kt_order, - index=1, - ) - # Index 2: QK output layout (128x16) - xegpu.set_anchor_layout( - qk_dpas_op, - sg_layout=layout_128x16_sg_layout, - sg_data=layout_128x16_sg_data, - inst_data=layout_128x16_inst_data, - index=2, - ) + # Layouts for the Q@K^T dpas: + qk_dpas_op = dpas_ops[0] + # Index 0: Q layout + xegpu.set_anchor_layout( + qk_dpas_op, + sg_layout=q_sg_layout, + sg_data=q_sg_data, + index=0, + ) + # Index 1: K^T layout + xegpu.set_anchor_layout( + qk_dpas_op, + sg_layout=kv_sg_layout, + sg_data=kt_sg_data, + index=1, + ) + # Index 2: QK output layout + xegpu.set_anchor_layout( + qk_dpas_op, + sg_layout=qk_sg_layout, + sg_data=qk_sg_data, + index=2, + ) - # Layouts for second 4 dpas ops (P@V): - for pv_dpas_op in dpas_ops[4:]: - # Index 0: QK (attention weights) layout - xegpu.set_anchor_layout( - pv_dpas_op, - sg_layout=qk_sg_layout, - sg_data=qk_sg_data, - inst_data=qk_inst_data, - index=0, - ) - # Index 1: V layout - xegpu.set_anchor_layout( - pv_dpas_op, - sg_layout=v_sg_layout, - sg_data=v_sg_data, - inst_data=v_inst_data, - index=1, - ) - # Index 2: Output layout - xegpu.set_anchor_layout( - pv_dpas_op, - sg_layout=out_sg_layout, - sg_data=out_sg_data, - inst_data=out_inst_data, - index=2, - ) + # Layouts for the P@V dpas: + pv_dpas_op = dpas_ops[1] + # Index 0: QK (attention weights) layout + xegpu.set_anchor_layout( + pv_dpas_op, + sg_layout=qk_sg_layout, + sg_data=qk_sg_data, + index=0, + ) + # Index 1: V layout + xegpu.set_anchor_layout( + pv_dpas_op, + sg_layout=kv_sg_layout, + sg_data=v_sg_data, + index=1, + ) + # Index 2: Output layout + xegpu.set_anchor_layout( + pv_dpas_op, + sg_layout=out_sg_layout, + sg_data=out_sg_data, + index=2, + ) if stop_at_stage == "xegpu-wg": raise PipelineInterrupt() From 0cd22b6dc628d26c0be502c82bb0b98e08151879 Mon Sep 17 00:00:00 2001 From: Charitha Saumya Date: Thu, 13 Aug 2026 18:38:59 +0000 Subject: [PATCH 3/3] match loads inside scf.for for prefetch --- lighthouse/schedule/xegpu/fused_attention_schedule.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/lighthouse/schedule/xegpu/fused_attention_schedule.py b/lighthouse/schedule/xegpu/fused_attention_schedule.py index f79442bf..82a3ca75 100644 --- a/lighthouse/schedule/xegpu/fused_attention_schedule.py +++ b/lighthouse/schedule/xegpu/fused_attention_schedule.py @@ -332,11 +332,9 @@ def bundle_xegpu_fused_attention_schedule( # the other wg-level layouts. nb_prefetch = parameters.get("nb_prefetch", 1) if nb_prefetch > 0: - # 3 load_nd ops: Q (hoisted out of the loop), then K and V inside it. - initial_load_nd_ops = match_and_split( - gpu_func, ops={"xegpu.load_nd"}, nhandles=3 - ) - for load_op in initial_load_nd_ops[1:]: + reduction_loop = match(gpu_func, ops={"scf.for"}) + kv_load_ops = match_and_split(reduction_loop, ops={"xegpu.load_nd"}, nhandles=2) + for load_op in kv_load_ops: xegpu.insert_prefetch(load_op, nb_prefetch=nb_prefetch) transform.apply_cse(gpu_func) canonicalize(gpu_func)