Xe SDPA Forward Mainloop
Key Loops + Shapes

Scope: simplified non-KV-cache and non-paged path. The kernel streams K/V blocks, keeps score/probability tiles in register fragments, and updates online softmax state across K blocks.

Q: [Sq, D] K: [Sk, D] V storage: [Dv, Sk] S / P tile: [BQ, BK] O split: [BQ, BV]

Accuracy notes

Scope

The loop and shape sections intentionally use the non-cache, non-paged mental model. Optional cache and paged-KV fields are only shown in the whole-file skeleton.

Dispatch policy

The actual policy tag is XeDefault<Stages>. VTiles is a separate template parameter of the mainloop specialization.

QK notation

The source comment says S = K * Q. The diagrams use the logical attention-shape view Q [BQ,BD] × Kᵀ [BD,BK] → S [BQ,BK].

Whole-file high-level skeleton

This map shows how the header is organized. The static file structure is shown first; the runtime skeleton with key source-code excerpts is shown after it.

Static structure of `xe_sdpa_fwd_mainloop.hpp`

Includes + namespacesPulls in CUTLASS/CUTE GEMM, subgroup, and MMA helpers, then opens the SDPA and FMHA collective namespaces.
Dispatch policy: XeDefault<Stages>Compile-time tag selecting the Intel Xe SDPA forward mainloop implementation.
Generic fallback templateSDPAFwdMainloop<...> default version only contains a compile-time “no specialization found” assertion.
XeDefault specializationThe real implementation. Most of the file lives inside this specialization.
1 Type aliases + compile-time flagsTile shapes, copy policies, MMA layouts, tensor element types, CausalMask, CachedKV, PagedKV, and VTiles.
2 Fragment / accumulator typesFragS for QK score tile, FragA for PV output accumulator, row fragments for online softmax.
3 Runtime args / paramsscale, optional mask, optional page-table metadata. to_underlying_arguments() converts scale for exp2.
4 Helperscan_implement(), optional get_physical_k_tile() for paged KV, and constructor wiring.
5 operator()Main execution path: builds tiled tensor views, copy/MMA objects, register fragments, prefetch objects, and launches K-block loops.
mainloop_body lambdaProcesses one K block: QK → V prefetch → mask → softmax → PV → future-K prefetch.
6 softmax() helperBlocked online softmax: row max, rescale, exp2, row sum update.

Runtime skeleton + key source code

Right-side snippets are shortened study excerpts. They show where each high-level skeleton piece appears in the file, while omitting cache/paged branches and some template details.

Dispatch policy + fallback

Policy tag selects the Xe implementation. Generic template fails if no matching specialization exists.

Source skeletoncompile-time dispatch

template <int Stages>
class XeDefault {};

template <class DispatchPolicy, ...>
struct SDPAFwdMainloop {
  static_assert(...,
    "Could not find a mainloop specialization.");
};

XeDefault specialization

The real implementation is the large specialization parameterized by masks, cache mode, MMA policies, tensor types, copy policies, and VTiles.

Source skeletonmain struct

template <int Stages, bool FullMask_, bool CausalMask_,
          bool CachedKV_, bool PagedKV_, class TiledMMAQK_, ...>
struct SDPAFwdMainloop<
    sdpa::XeDefault<Stages>,
    CausalMask_, FullMask_, CachedKV_, PagedKV_,
    TiledMMAQK_, TiledMMAPV_, VTiles_, ...> {
  ...
};

Type aliases + fragments

Compile-time tile shapes and register fragments are defined before runtime execution.

Source skeletontypes / fragments

using TileShapeQK = decltype(TiledMMAQK{}.tile_mnk());
using TileShapePV = decltype(TiledMMAPV{}.tile_mnk());

using FragS    = ...;  // QK score tile
using FragSRow = ...;  // row max/sum state
using FragA    = ...;  // PV output accumulator

Arguments / params

Runtime arguments carry scale, optional full mask, and optional paged-KV metadata. Scale is converted for exp2.

Source skeletonruntime args

struct Arguments {
  ElementS const scale;
  ElementM const* mask = nullptr;
  int const* ptr_page_table = nullptr;
  int page_size = 0;
  int const* num_pages_per_seq = nullptr;
};

static Params to_underlying_arguments(Arguments const& args) {
  Params params = args;
  params.scale *= log2(e);   // softmax uses exp2()
  return params;
}

`operator()` setup

Builds coordinate tensors, local tiles, copy objects, MMA objects, and thread-local fragments.

Source skeletonoperator setup

Tensor cQ = make_identity_tensor(Q_2D.shape());
Tensor cK = make_identity_tensor(K_2D.shape());
Tensor cV = make_identity_tensor(V_2D.shape());
Tensor cP = make_identity_tensor(...);  // (q,k)

Tensor gQ = local_tile(Q_2D, tile_shape_q, ...);
Tensor gK = local_tile(K_2D, tile_shape_k, ...);
Tensor gV_split = local_tile(V_2D, tile_shape_v, ...);

TiledCopyQ copy_q{Q_2D};
TiledCopyK copy_k{K_2D};
TiledCopyV copy_v{V_2D};
TiledMMAQK mma_qk{};
TiledMMAPV mma_pv{};

Startup prefetch + init

Prefetch Q and first Stages K blocks, then initialize O accumulator and online-softmax row state for the first block.

Source skeletonbefore K loop

int kblocks_cache = ceil_div(seq_len_kv_cache, get<1>(TileShapeQK{}));

prefetch(prefetch_q, pQgQ(...));

for (int D = 0; D < size<4>(pKgK); D++) {
  for (int K = 0; K < Stages; K++) {
    // non-cache path: kblocks_cache == 0, so this is normal K
    prefetch(prefetch_k, pKgK(_, _, _, K, D));
  }
}

if (blk_k0 == 0) {
  clear(tArA);
  fill(tA_max, -INFINITY);
  clear(tA_sum);
}

`mainloop_body(K)`

One call handles one K block: QK, V prefetch, masks, softmax, PV, and future-K prefetch.

Source skeletonone K block

auto mainloop_body = [&](auto is_cache, int K, ...) {
  /* GEMM 1: S = K * Q */
  // logical shape view: Q[BQ,BD] x K^T[BD,BK]
  clear(tSrS);
  for (int D = 0; D < size<4>(tKgK); D++) {
    copy(... Q ..., tQrQ);
    copy(... K ..., tKrK);
    reorder(tQrQ, tSrQ);
    reorder(tKrK, tSrK);
    cute::gemm(mma_qk, tSrQ, tSrK, tSrS);
  }

  // mask + softmax
  ...
  auto rescale = softmax(K == blk_k0, tSrS, tA_max, tA_sum);
  reorder(tSrS, tArP);

  // GEMM 2: A += P * V
  for (int VV = 0; VV < VTiles; VV++) {
    copy(... V ..., tVrV);
    reorder(tVrV, tArV);
    cute::gemm(mma_pv, tArP, tArV, tArA(_,_,_,VV));
  }
};

Outer K-block loop

Ignoring cache/paged mode, this loop repeatedly calls mainloop_body for normal K blocks.

Source skeletonstream K blocks

for (int K = max(blk_k0, kblocks_cache);
     K < blk_k1; K++) {
  mainloop_body(cute::false_type{}, K,
                copy_k, copy_v,
                ...);
}

`softmax()` helper

Updates online row max/sum and returns rescale [BQ] for previous O accumulator correction.

Source skeletononline softmax

FragSRow softmax(bool first_block,
                 FragS& tS,
                 FragSRow& tS_max,
                 FragSRow& tS_sum) {
  auto tS_bmax = reduce<1>(tS, sycl::maximum{});

  // new max, rescale old max/sum
  tS_max(i) = max(tS_max(i), params.scale * tS_bmax(i));
  rescale(i) = exp2(old_max - tS_max(i));

  // convert score tile into exp/softmax numerator
  tS(i) = sycl::native::exp2(params.scale * tS(i) - tS_max(...));

  auto tS_bsum = reduce<1>(tS, sycl::plus{});
  if (!first_block) tS_sum(i) *= rescale(i);
  tS_sum(i) += tS_bsum(i);
  return rescale;
}

policy specialization types operator() mainloop_body(K) softmax() outer K loop

SageV1 concrete instance: head_dim=128, seq_q=512, seq_k=512

This section connects the launcher dispatch, instantiated tile shapes, and actual loop counts for the SageV1 prefill path. Assumption: non-varlen, non-cache, non-paged prefill with head_size_qk = head_size_vo = 128.

Who decides the instantiated mainloop?

1. Public entrysage_prefill(...) builds Options, including dtype, PV dtype, head_dim, causal flag, scale pointers, and use_int8_pv.

2. Launcher selectionselect_sage_prefill_launcher(q_dtype, pv_dtype, head_dim, use_int8_pv) chooses the 64- or 128-head-dim launcher. For head_dim=128, it selects a *_128_sage launcher.

3. Tile-shape definitionlaunch_sage_prefill_kernel_128 defines ShapeQK, ShapePV, ShapeOut, SubgroupLayoutQK, and PipelineStages.

4. SageConfigSageConfig builds TiledMMAQK / TiledMMAPV, derives VTiles, and instantiates SAGEV1FwdMainloop.

5. Runtime variantSageConfig::run(options) selects varlen/cache/paged/scheduler template parameters. In this example: isVarLen=false, CachedKV=false, PagedKV=false.

Key source skeleton

KernelLauncher launcher =
  select_sage_prefill_launcher(q_dtype, pv_dtype,
                               head_dim, use_int8_pv);

inline int launch_sage_prefill_kernel_128(...) {
  constexpr int PipelineStages = 2;

  using ShapeQK  = Shape<_256, _64, _32>;
  using ShapePV  = Shape<_256, _32, _64>;
  using ShapeOut = Shape<_256, _128>;
  using SubgroupLayoutQK = Layout<Shape<_16, _1, _1>>;

  return options.is_causal
    ? SageConfig<true,  ..., ShapeQK, ShapePV, ShapeOut, ...>::run(options)
    : SageConfig<false, ..., ShapeQK, ShapePV, ShapeOut, ...>::run(options);
}

constexpr int VTiles =
  get<1>(ShapeOut{}) / get<1>(ShapePV{});  // 128 / 32 = 4

using MainloopDispatchPolicy =
  cutlass::sage::XeDefault<PipelineStages>;

using CollectiveMainloop =
  SAGEV1FwdMainloop<MainloopDispatchPolicy, ..., TiledMMAQK,
                    TiledMMAPV, VTiles, ...>;

QK tile shape

ShapeQK = Shape<256, 64, 32>

GEMM order: M,N,K. Here M=256 query rows, N=64 key-token columns, and K=32 head-dim reduction chunk.

PV tile shape

ShapePV = Shape<256, 32, 64>

GEMM order: M,N,K. Here M=256 query rows, N=32 output/value-dim split, and K=64 key-token reduction.

Output tile shape

ShapeOut = Shape<256, 128>

One output tile covers 256 query rows and the full 128 V/O head dimension. Since PV covers only 32 output dims each time, VTiles = 4.

Shape interpretation

Object	Meaning	Concrete operation	Result
`ShapeQK<256,64,32>`	QK GEMM tile: `M=query`, `N=key`, `K=head dim chunk`	`Q [256,32] × Kᵀ [32,64]`	`S [256,64]`
`ShapePV<256,32,64>`	PV GEMM tile: `M=query`, `N=output dim split`, `K=key`	`P [256,64] × V [64,32]`	`O_split [256,32]`
`ShapeOut<256,128>`	Full output tile for one Q block	`4 × O_split [256,32]`	`O [256,128]`
`PipelineStages = 2`	K prefetch pipeline depth	Startup prefetches `K0,K1`; processing `Kx` prefetches `Kx+2`	K pipeline stays two blocks ahead when in range

Loop counts for seq_q=512, seq_k=512, head_dim=128

Q blocks
512 / 256

K blocks
512 / 64

D-loop chunks per K block
128 / 32

PV/VV loops per K block
128 / 32

One Q block timeline

Q tile

Q [256,128]

K stream

K0 0:64 K1 64:128 K2 128:192 K3 192:256 K4 256:320 K5 320:384 K6 384:448 K7 448:512

Per K block

4 × QK D-loop S/P [256,64] 4 × PV VV-loop O update [256,128]

Prefetch

startup: K0,K1 K0 → prefetch K2 K1 → prefetch K3 ... K5 → prefetch K7

Operation counts

Scope	QK MMA tiles	PV MMA tiles	Softmax tiles
One K block inside one Q block	`4`	`4`	`1 × [256,64]`
One Q block	`8 K blocks × 4 = 32`	`8 K blocks × 4 = 32`	`8`
Full sequence per batch/head	`2 Q blocks × 32 = 64`	`2 Q blocks × 32 = 64`	`2 × 8 = 16`

These counts are conceptual tile-level counts for the selected mainloop shape. Actual issued hardware instructions are further determined by TiledMMAHelper, MMA atom, subgroup layout, and compiler lowering.

Compact summary

SageV1 prefill, head_dim=128, seq_q=512, seq_k=512:

Selected launcher:
  launch_sage_prefill_kernel_128

Instantiated compile-time shapes:
  ShapeQK  = Shape<256, 64, 32>   // QK GEMM: M=query, N=key, K=head-dim chunk
  ShapePV  = Shape<256, 32, 64>   // PV GEMM: M=query, N=value/output split, K=key
  ShapeOut = Shape<256, 128>      // output tile
  Stages   = 2
  VTiles   = 128 / 32 = 4

Loop counts:
  Q blocks = 512 / 256 = 2
  K blocks = 512 / 64  = 8
  D loops  = 128 / 32  = 4 per K block
  VV loops = 128 / 32  = 4 per K block

Per Q block:
  8 K blocks
  32 QK MMA tiles
  32 PV MMA tiles
  8 softmax tiles [256,64]

Full sequence per batch/head:
  64 QK MMA tiles
  64 PV MMA tiles
  16 softmax tiles [256,64]

Shape legend

Sequence dims

Sq: query length
Sk: key/value length

Head dims

D: Q/K head dim
Dv: V/output head dim

Tile dims

BQ: Q tile rows
BK: K tile rows
BD: inner D slice
BV: V split width

Loop knobs

Stages: prefetched future K blocks
VTiles: V/output head-dim splits

One outer K-block iteration

Q [BQ, D] × Kᵀ [D, BK] → S/P [BQ, BK] × V [BK, BV] → O [BQ, BV]

The full attention matrix is not materialized. Each K block produces one score/probability tile [BQ, BK], which is immediately consumed by the PV GEMM.

Interactive prefetch view

The startup prefetch covers the first Stages K blocks. For each prefetched K block, it walks all D/head-dim slices.

Stages 3

D tiles 4

VTiles 4

K timeline rows 5

Startup K prefetch

Each small D cell represents one [BK, BD] K slice. A whole K block is [BK, D].

Key loop pipeline

Read this as

compute current K block

V prefetch / PV current K block's V

prefetch future K block K + Stages

Assume kblocks_cache = 0. So the non-cache loop directly processes normal K blocks.

Mainloop body with shapes + key source code

Left side shows the logical shape transition. Right side shows the corresponding shortened source-code region from mainloop_body or softmax(). Ellipses mean cache/paged branches or bounds details are omitted for readability.

QK D-loop

Loop over Q/K head-dim slices.

Q [BQ, BD] × Kᵀ [BD, BK] → S [BQ, BK]

Logical attention view is Q × Kᵀ. The source comment says S = K * Q because of the MMA operand/layout convention. Repeated over all D tiles; each partial MMA accumulates into tSrS.

Source: GEMM 1shape: [BQ,BD]×[BD,BK]

/* GEMM 1: S = K * Q */
clear(tSrS);
for (int D = 0; D < size<4>(tKgK); D++) {
  copy(copy_q, tQgQ(_, _, _, D), tQrQ);
  copy(copy_k_cur, tKgK_cur(_, _, _, k_idx, D), tKrK);
  reorder(tQrQ, tSrQ);
  reorder(tKrK, tSrK);
  cute::gemm(mma_qk, tSrQ, tSrK, tSrS);
}

V prefetch VV-loop

Before PV GEMM, prefetch V splits for the current K block.

V split storage [BV, BK]

Math view is effectively [BK, BV] after MMA-layout reorder.

Source: V prefetchshape: per VV [BV,BK]

for (int VV = 0; VV < VTiles; VV++) {
  prefetch(prefetch_v_cur,
           pVgV_cur(_, _, _, VV, k_idx));
}

Mask

Elementwise update of score tile.

S [BQ, BK] → S [BQ, BK]

Depending on compile-time flags and boundary conditions, causal, full, or remainder masking changes score values but not the logical tile shape.

Source: mask branchesshape unchanged

if constexpr (!is_cache && CausalMask) {
  ...
  tSrS(i) = ElementS(-INFINITY);
}

if (check_remainder_k && K == total_blk - 1) {
  ...
  tSrS(i) = sycl::fmin(tSrS(i), ...);
}

Online softmax

Reduce over K columns, then broadcast row state back to the score tile.

S [BQ, BK] → row_max [BQ] + row_sum [BQ] → P [BQ, BK]

softmax() returns rescale [BQ]; then tSrS is reordered into the PV A-operand layout tArP.

Source: softmax + reorder[BQ,BK] → [BQ]

auto rescale = softmax(K == blk_k0,
                       tSrS, tA_max, tA_sum);
reorder(tSrS, tArP);

// inside softmax()
auto tS_bmax = reduce<1>(tS, sycl::maximum{});
tS(i) = sycl::native::exp2(params.scale * tS(i) - ...);
auto tS_bsum = reduce<1>(tS, sycl::plus{});
tS_sum(i) += tS_bsum(i);

PV VV-loop

Loop over V/output splits.

P [BQ, BK] × V [BK, BV] → O_split [BQ, BV]

If this is not the first K block, old O is row-rescaled by rescale [BQ] before adding the current PV result.

Source: GEMM 2[BQ,BK]×[BK,BV]

for (int VV = 0; VV < VTiles; VV++) {
  copy(copy_v_cur, tVgV_cur(_, _, _, VV, k_idx), tVrV);
  reorder(tVrV, tArV);

  if (K != blk_k0)
    tArA(_, _, _, VV)(i) *= broadcast<0>(rescale, ...);

  cute::gemm(mma_pv, tArP, tArV, tArA(_, _, _, VV));
}

Future-K prefetch D-loop

After processing current K, request the future K block.

K_next = K + Stages → all [BK, BD] slices

This keeps the K pipeline ahead of compute.

Source: K prefetchfuture [BK,D]

int K_next = K + Stages;
for (int D = 0; D < size<4>(pKgK); D++) {
  prefetch(prefetch_k,
           pKgK(_, _, _, K_next - kblocks_cache, D));
}

Full simplified pseudocode

// Global logical tensors
// Q: [Sq, D], K: [Sk, D], V storage: [Dv, Sk], O: [Sq, Dv]

// Startup
prefetch Q tile                         // [BQ, D] as D slices [BQ, BD]
for D_tile in all_D_tiles:
  for k_prefetch in 0 .. Stages-1:
    prefetch K[k_prefetch, D_tile]       // [BK, BD]

if first_K_block:
  O_acc   = 0                            // per split: [BQ, BV]
  row_max = -inf                         // [BQ]
  row_sum = 0                            // [BQ]

// Main non-cache K loop
for K_block in K_blocks:

  // 1. QK GEMM
  S = 0                                  // [BQ, BK]
  for D_tile in all_D_tiles:
    Q_frag = Q[:, D_tile]                // [BQ, BD]
    K_frag = K[K_block, D_tile]          // [BK, BD]
    S += Q_frag @ K_frag.T               // [BQ, BK]

  // 2. V prefetch
  for VV in VTiles:
    prefetch V[VV, K_block]              // storage [BV, BK]

  // 3. Mask
  S = apply_mask(S)                      // [BQ, BK]

  // 4. Online softmax
  block_max = max(S, axis=K)             // [BQ]
  new_max   = max(row_max, block_max)    // [BQ]
  rescale   = exp(old_max - new_max)     // [BQ]
  P         = exp(S - new_max[:, None])  // [BQ, BK]
  row_sum   = row_sum * rescale + sum(P, axis=K) // [BQ]
  row_max   = new_max                    // [BQ]

  // 5. PV GEMM, split along V/output dimension
  for VV in VTiles:
    V_frag = V[VV, K_block]              // storage [BV, BK], math [BK, BV]
    O_acc[VV] *= rescale[:, None]        // [BQ, BV]
    O_acc[VV] += P @ V_frag.T            // [BQ, BV]

  // 6. Future-K prefetch
  for D_tile in all_D_tiles:
    prefetch K[K_block + Stages, D_tile] // [BK, BD]

Compact shape table

Loop / operation	Input shape	Output / state shape	Purpose
Startup Q prefetch	Q tile `[BQ, D]`	D slices `[BQ, BD]`	Bring query tile ahead of QK compute.
Startup K prefetch	first `Stages` K blocks	each block split as `[BK, BD]`	Warm K pipeline before mainloop starts.
QK D-loop	`[BQ, BD] × [BD, BK]`	`S [BQ, BK]`	Accumulate score tile over D.
Mask	`S [BQ, BK]`	`S [BQ, BK]`	Set invalid scores to negative infinity, apply a full mask when that path is enabled, or handle the last partial K tile.
Softmax	`S [BQ, BK]`	`P [BQ, BK]`, `row_max/sum [BQ]`	Online row-wise stable softmax.
PV VV-loop	`P [BQ, BK] × V [BK, BV]`	`O_split [BQ, BV]`	Accumulate output split.
Future K prefetch	`K_block + Stages`	all D slices `[BK, BD]`	Keep next K blocks ready for later iterations.

Mental model

Prefetch window

Before compute starts: request first Stages K blocks, each across all D slices.

→

Streaming score tile

For one K block, build S [BQ, BK] by looping over D slices.

→

Consume immediately

Softmax turns S into P, then P @ V updates output. No full attention matrix is stored.

Source inspected: Intel AutoRound xe_sdpa_fwd_mainloop.hpp. This visualization intentionally ignores KV cache and paged attention branches and explains the normal K/V path only.

Xe SDPA Forward MainloopKey Loops + Shapes