diff --git a/include/common/base_types.dp.hpp b/include/common/base_types.dp.hpp index 6e25529..eceb163 100644 --- a/include/common/base_types.dp.hpp +++ b/include/common/base_types.dp.hpp @@ -477,7 +477,10 @@ template <> struct convertor { }; template <> struct convertor { static inline bf16_2 convert(const sycl::float2 &u) { - return sycl::vec(u.x(), u.y()); + // Reuse the scalar converter: the bf16_2 vec constructor lowers to a + // software RNE sequence (shr/and/add3 + NaN branch) on Xe2. + return bf16_2(convertor::convert(u.x()), + convertor::convert(u.y())); } }; template <> struct convertor { diff --git a/include/ops/warp/memory/tile/gl_desc.dp.hpp b/include/ops/warp/memory/tile/gl_desc.dp.hpp index 68ccf1f..6d747f8 100644 --- a/include/ops/warp/memory/tile/gl_desc.dp.hpp +++ b/include/ops/warp/memory/tile/gl_desc.dp.hpp @@ -37,8 +37,9 @@ * - Register cost: each live `gl_desc` is ~1 GRF. GEMM runs near the 256-GRF * ceiling — share one descriptor per operand-stream and advance X/Y rather * than holding many; re-run spillcheck after adopting. - * - The `.load()` inline asm is unvalidated in THIS tree until you build+run - * it on PVC. Confirm bit-exact against the flat load before trusting it. + * - Validated on B70 (Xe2) by the attention forward K/V loads (transposed, + * VNNI and prefetch), which match the flat path's output. PVC builds but + * has not been run; confirm against the flat load there before trusting it. * ============================================================================ */ @@ -60,20 +61,26 @@ namespace kittens { * @tparam VNNI true selects the transform/VNNI read (`nt`, matches * Subgroup2DBlockLoadTransformINTEL — the B operand); * false selects the plain read (`nn` — the A operand). + * @tparam Transpose true selects the transposed read (`tn`, matches + * Subgroup2DBlockLoadTransposeINTEL; NumBlocks must be 1). * * Coordinates are ELEMENT coordinates (identical to the {col,row} the flat * intrinsic takes): X = contiguous/column element index, Y = strided/row * element index. */ template + bool VNNI = false, bool Transpose = false> struct gl_desc { + static_assert(!(VNNI && Transpose), "VNNI and transposed reads are exclusive"); + static_assert(!Transpose || NumBlocks == 1, "transposed reads take one block"); + int *payload = nullptr; ///< the block2d address payload (one GRF, built once) static constexpr int sg_size = 16; static constexpr int atom_width = BlockWidth * NumBlocks; ///< total contiguous elems // Per-lane storage vector matching the flat load's deposited bytes. - using elem_t = std::conditional_t<(Bits <= 16), short, int>; + // VNNI packs 32/Bits rows into each dword, so it is read back as int. + using elem_t = std::conditional_t<(Bits <= 16 && !VNNI), short, int>; static constexpr int lane_count = (atom_width * BlockHeight * Bits / sg_size) / static_cast(sizeof(elem_t) * 8); @@ -92,14 +99,18 @@ struct gl_desc { template inline void init(const GL &g) { using U = typename GL::dtype; - const long base = reinterpret_cast(g.raw_ptr); - const int width_bytes = static_cast(g.template shape<3>() * sizeof(U)); - const int height = static_cast(g.template shape<2>()); - const int pitch_bytes = - static_cast(g.template stride<2>() * sizeof(U)); + init(g.raw_ptr, static_cast(g.template shape<3>() * sizeof(U)), + static_cast(g.template shape<2>()), + static_cast(g.template stride<2>() * sizeof(U))); + } + + /// Build the payload for an arbitrary 2D surface (sizes in bytes / rows). + inline void init(const void *base, int width_bytes, int height, + int pitch_bytes) { payload = __builtin_IB_subgroup_createBlock2DAddressPayload( - base, width_bytes - 1, height - 1, pitch_bytes - 1, - /*blockX=*/0, /*blockY=*/0, BlockWidth, BlockHeight, NumBlocks); + reinterpret_cast(base), width_bytes - 1, height - 1, + pitch_bytes - 1, /*blockX=*/0, /*blockY=*/0, BlockWidth, BlockHeight, + NumBlocks); } /// Set both surface coordinates (element units). @@ -127,22 +138,29 @@ struct gl_desc { /** * @brief Load the block(s) into register-tile storage. `dst` is the same * pointer the flat path writes (e.g. ®.tiles[i][j].data[0]); the - * deposited bytes are bit-identical. Mirrors sycl-tla XE_LOAD_2D / - * XE_LOAD_2D_VNNI. + * deposited bytes are bit-identical. DX/DY are compile-time offsets + * (in `Bits`-sized elements / rows) added to the current X/Y, so a + * fixed pattern of loads around one coordinate needs no payload + * update. Mirrors sycl-tla XE_LOAD_2D / _VNNI / _TRANSPOSE. */ - template inline void load(T *dst) const { + template inline void load(T *dst) const { #ifdef __SYCL_DEVICE_ONLY__ auto &dv = *reinterpret_cast(dst); - if constexpr (VNNI) { - asm("lsc_load_block2d.ugm (M1, 1) %0:d%2.%3x%4x%5nt flat[%1+(0,0)]" + if constexpr (Transpose) { + asm("lsc_load_block2d.ugm (M1, 1) %0:d%2.%3x%4tn flat[%1+(%5,%6)]" + : "=rw"(dv) + : "rw.u"(payload), "P"(Bits), "P"(BlockWidth), "P"(BlockHeight), + "P"(DX), "P"(DY)); + } else if constexpr (VNNI) { + asm("lsc_load_block2d.ugm (M1, 1) %0:d%2.%3x%4x%5nt flat[%1+(%6,%7)]" : "=rw"(dv) : "rw.u"(payload), "P"(Bits), "P"(NumBlocks), "P"(BlockWidth), - "P"(BlockHeight)); + "P"(BlockHeight), "P"(DX), "P"(DY)); } else { - asm("lsc_load_block2d.ugm (M1, 1) %0:d%2.%3x%4x%5nn flat[%1+(0,0)]" + asm("lsc_load_block2d.ugm (M1, 1) %0:d%2.%3x%4x%5nn flat[%1+(%6,%7)]" : "=rw"(dv) : "rw.u"(payload), "P"(Bits), "P"(NumBlocks), "P"(BlockWidth), - "P"(BlockHeight)); + "P"(BlockHeight), "P"(DX), "P"(DY)); } #else (void)dst; @@ -162,10 +180,11 @@ struct gl_desc { * load coordinate is clobbered. See the migration recipe. Each extra * descriptor is ~1 GRF; re-run spillcheck. */ - inline void prefetch() const { + template inline void prefetch() const { #ifdef __SYCL_DEVICE_ONLY__ - asm("lsc_load_block2d.ugm.ca.ca (M1, 1) %%null:d%1.%2x%3nn flat[%0+(0,0)]" - :: "rw.u"(payload), "P"(Bits), "P"(atom_width), "P"(BlockHeight)); + asm("lsc_load_block2d.ugm.ca.ca (M1, 1) %%null:d%1.%2x%3nn flat[%0+(%4,%5)]" + :: "rw.u"(payload), "P"(Bits), "P"(atom_width), "P"(BlockHeight), + "P"(DX), "P"(DY)); #endif } }; diff --git a/include/ops/warp/register/tile/reductions.dp.hpp b/include/ops/warp/register/tile/reductions.dp.hpp index 5e6f01f..4d5ed32 100644 --- a/include/ops/warp/register/tile/reductions.dp.hpp +++ b/include/ops/warp/register/tile/reductions.dp.hpp @@ -37,6 +37,154 @@ inline float reduce_16x16(const sycl::float2 *first_half, * @param[in] src The source matrix on which to perform the reduction. * @param[in] src_accum The initial value of the accumulator, used when reset is false. */ +#ifdef __SYCL_DEVICE_ONLY__ +// Register-region transpose-reduce of 16 (or 8) per-lane row partials across the +// 16 SIMD lanes: lane L gets the reduction of row L. Every stage is a plain SIMD16 +// mov/op on GRF regions, so there are no a0-indirect shuffles (which serialize on +// Xe). Same data movement as sycl-tla's hreduce16/8_float_* (BSD-3-Clause, +// Copyright 2025 Intel Corporation). For 8 rows, lanes 8-15 are left undefined. +typedef float sk_float16v __attribute__((ext_vector_type(16))); +typedef float sk_float8v __attribute__((ext_vector_type(8))); +#define SK_DEFINE_HREDUCE16(op) \ + static inline float sk_hreduce16_##op(sk_float16v x) { \ + float y; \ + asm("{\n" \ + ".decl SKR_S v_type=G type=F num_elts=512 align=wordx32\n" \ + ".decl SKR_T v_type=G type=F num_elts=256 align=wordx32\n" \ + "mov (M1_NM,16) SKR_S(0,0)<1> %1(0,0)<16;8,1>\n" \ + "mov (M1_NM,16) SKR_S(1,0)<1> %1(0,8)<16;8,1>\n" \ + "mov (M1_NM,16) SKR_S(2,0)<1> %1(2,0)<16;8,1>\n" \ + "mov (M1_NM,16) SKR_S(3,0)<1> %1(2,8)<16;8,1>\n" \ + "mov (M1_NM,16) SKR_S(4,0)<1> %1(4,0)<16;8,1>\n" \ + "mov (M1_NM,16) SKR_S(5,0)<1> %1(4,8)<16;8,1>\n" \ + "mov (M1_NM,16) SKR_S(6,0)<1> %1(6,0)<16;8,1>\n" \ + "mov (M1_NM,16) SKR_S(7,0)<1> %1(6,8)<16;8,1>\n" \ + "mov (M1_NM,16) SKR_S(8,0)<1> %1(8,0)<16;8,1>\n" \ + "mov (M1_NM,16) SKR_S(9,0)<1> %1(8,8)<16;8,1>\n" \ + "mov (M1_NM,16) SKR_S(10,0)<1> %1(10,0)<16;8,1>\n" \ + "mov (M1_NM,16) SKR_S(11,0)<1> %1(10,8)<16;8,1>\n" \ + "mov (M1_NM,16) SKR_S(12,0)<1> %1(12,0)<16;8,1>\n" \ + "mov (M1_NM,16) SKR_S(13,0)<1> %1(12,8)<16;8,1>\n" \ + "mov (M1_NM,16) SKR_S(14,0)<1> %1(14,0)<16;8,1>\n" \ + "mov (M1_NM,16) SKR_S(15,0)<1> %1(14,8)<16;8,1>\n" \ + #op " (M1_NM,16) SKR_T(0,0)<1> SKR_S(0,0)<1;1,0> SKR_S(1,0)<1;1,0>\n" \ + #op " (M1_NM,16) SKR_T(1,0)<1> SKR_S(2,0)<1;1,0> SKR_S(3,0)<1;1,0>\n" \ + #op " (M1_NM,16) SKR_T(2,0)<1> SKR_S(4,0)<1;1,0> SKR_S(5,0)<1;1,0>\n" \ + #op " (M1_NM,16) SKR_T(3,0)<1> SKR_S(6,0)<1;1,0> SKR_S(7,0)<1;1,0>\n" \ + #op " (M1_NM,16) SKR_T(4,0)<1> SKR_S(8,0)<1;1,0> SKR_S(9,0)<1;1,0>\n" \ + #op " (M1_NM,16) SKR_T(5,0)<1> SKR_S(10,0)<1;1,0> SKR_S(11,0)<1;1,0>\n" \ + #op " (M1_NM,16) SKR_T(6,0)<1> SKR_S(12,0)<1;1,0> SKR_S(13,0)<1;1,0>\n" \ + #op " (M1_NM,16) SKR_T(7,0)<1> SKR_S(14,0)<1;1,0> SKR_S(15,0)<1;1,0>\n" \ + "mov (M1_NM,16) SKR_S(16,0)<1> SKR_T(0,0)<8;4,1>\n" \ + "mov (M1_NM,16) SKR_S(17,0)<1> SKR_T(0,4)<8;4,1>\n" \ + "mov (M1_NM,16) SKR_S(18,0)<1> SKR_T(2,0)<8;4,1>\n" \ + "mov (M1_NM,16) SKR_S(19,0)<1> SKR_T(2,4)<8;4,1>\n" \ + "mov (M1_NM,16) SKR_S(20,0)<1> SKR_T(4,0)<8;4,1>\n" \ + "mov (M1_NM,16) SKR_S(21,0)<1> SKR_T(4,4)<8;4,1>\n" \ + "mov (M1_NM,16) SKR_S(22,0)<1> SKR_T(6,0)<8;4,1>\n" \ + "mov (M1_NM,16) SKR_S(23,0)<1> SKR_T(6,4)<8;4,1>\n" \ + #op " (M1_NM,16) SKR_T(8,0)<1> SKR_S(16,0)<1;1,0> SKR_S(17,0)<1;1,0>\n" \ + #op " (M1_NM,16) SKR_T(9,0)<1> SKR_S(18,0)<1;1,0> SKR_S(19,0)<1;1,0>\n" \ + #op " (M1_NM,16) SKR_T(10,0)<1> SKR_S(20,0)<1;1,0> SKR_S(21,0)<1;1,0>\n" \ + #op " (M1_NM,16) SKR_T(11,0)<1> SKR_S(22,0)<1;1,0> SKR_S(23,0)<1;1,0>\n" \ + "mov (M1_NM,16) SKR_S(24,0)<1> SKR_T(8,0)<4;2,1>\n" \ + "mov (M1_NM,16) SKR_S(25,0)<1> SKR_T(8,2)<4;2,1>\n" \ + "mov (M1_NM,16) SKR_S(26,0)<1> SKR_T(10,0)<4;2,1>\n" \ + "mov (M1_NM,16) SKR_S(27,0)<1> SKR_T(10,2)<4;2,1>\n" \ + #op " (M1_NM,16) SKR_T(12,0)<1> SKR_S(24,0)<1;1,0> SKR_S(25,0)<1;1,0>\n" \ + #op " (M1_NM,16) SKR_T(13,0)<1> SKR_S(26,0)<1;1,0> SKR_S(27,0)<1;1,0>\n" \ + "mov (M1_NM,16) SKR_S(28,0)<1> SKR_T(12,0)<2;1,0>\n" \ + "mov (M1_NM,16) SKR_S(29,0)<1> SKR_T(12,1)<2;1,0>\n" \ + #op " (M1_NM,16) %0(0,0)<1> SKR_S(28,0)<1;1,0> SKR_S(29,0)<1;1,0>\n" \ + "}\n" : "=rw"(y) : "rw"(x)); \ + return y; \ + } + +#define SK_DEFINE_HREDUCE8(op) \ + static inline float sk_hreduce8_##op(sk_float8v x) { \ + float y; \ + asm("{\n" \ + ".decl SKR_S v_type=G type=F num_elts=256 align=wordx32\n" \ + ".decl SKR_T v_type=G type=F num_elts=128 align=wordx32\n" \ + "mov (M1_NM,16) SKR_S(0,0)<1> %1(0,0)<16;8,1>\n" \ + "mov (M1_NM,16) SKR_S(1,0)<1> %1(0,8)<16;8,1>\n" \ + "mov (M1_NM,16) SKR_S(2,0)<1> %1(2,0)<16;8,1>\n" \ + "mov (M1_NM,16) SKR_S(3,0)<1> %1(2,8)<16;8,1>\n" \ + "mov (M1_NM,16) SKR_S(4,0)<1> %1(4,0)<16;8,1>\n" \ + "mov (M1_NM,16) SKR_S(5,0)<1> %1(4,8)<16;8,1>\n" \ + "mov (M1_NM,16) SKR_S(6,0)<1> %1(6,0)<16;8,1>\n" \ + "mov (M1_NM,16) SKR_S(7,0)<1> %1(6,8)<16;8,1>\n" \ + #op " (M1_NM,16) SKR_T(0,0)<1> SKR_S(0,0)<1;1,0> SKR_S(1,0)<1;1,0>\n" \ + #op " (M1_NM,16) SKR_T(1,0)<1> SKR_S(2,0)<1;1,0> SKR_S(3,0)<1;1,0>\n" \ + #op " (M1_NM,16) SKR_T(2,0)<1> SKR_S(4,0)<1;1,0> SKR_S(5,0)<1;1,0>\n" \ + #op " (M1_NM,16) SKR_T(3,0)<1> SKR_S(6,0)<1;1,0> SKR_S(7,0)<1;1,0>\n" \ + "mov (M1_NM,16) SKR_S(8,0)<1> SKR_T(0,0)<8;4,1>\n" \ + "mov (M1_NM,16) SKR_S(9,0)<1> SKR_T(0,4)<8;4,1>\n" \ + "mov (M1_NM,16) SKR_S(10,0)<1> SKR_T(2,0)<8;4,1>\n" \ + "mov (M1_NM,16) SKR_S(11,0)<1> SKR_T(2,4)<8;4,1>\n" \ + #op " (M1_NM,16) SKR_T(4,0)<1> SKR_S(8,0)<1;1,0> SKR_S(9,0)<1;1,0>\n" \ + #op " (M1_NM,16) SKR_T(5,0)<1> SKR_S(10,0)<1;1,0> SKR_S(11,0)<1;1,0>\n" \ + "mov (M1_NM,16) SKR_S(12,0)<1> SKR_T(4,0)<4;2,1>\n" \ + "mov (M1_NM,16) SKR_S(13,0)<1> SKR_T(4,2)<4;2,1>\n" \ + #op " (M1_NM,16) SKR_T(6,0)<1> SKR_S(12,0)<1;1,0> SKR_S(13,0)<1;1,0>\n" \ + "mov (M1_NM,8) SKR_S(14,0)<1> SKR_T(6,0)<2;1,0>\n" \ + "mov (M1_NM,8) SKR_S(15,0)<1> SKR_T(6,1)<2;1,0>\n" \ + #op " (M1_NM,8) %0(0,0)<1> SKR_S(14,0)<1;1,0> SKR_S(15,0)<1;1,0>\n" \ + "}\n" : "=rw"(y) : "rw"(x)); \ + return y; \ + } +SK_DEFINE_HREDUCE16(max) +SK_DEFINE_HREDUCE16(add) +SK_DEFINE_HREDUCE8(max) +SK_DEFINE_HREDUCE8(add) +#undef SK_DEFINE_HREDUCE16 +#undef SK_DEFINE_HREDUCE8 +#endif + +/** + * @brief Transposing butterfly: reduce ROWS per-lane row partials across the 16 + * SIMD lanes so that lane L ends with the reduction of row L (L % ROWS). + * + * Each stage swaps half of the remaining rows with the partner lane, so all rows + * are reduced together: ROWS-1 shuffles (plus one fold for ROWS=8) instead of + * 4 per row. Shuffles lower to serialized a0-indirect moves on Xe, so this + * removes ~3/4 of them from softmax row reductions. + */ +template +static inline type transpose_reduce_rows(type (&a)[ROWS], int lane) { + static_assert(ROWS == 16 || ROWS == 8, "16 or 8 rows per SIMD16 tile"); +#ifdef __SYCL_DEVICE_ONLY__ + constexpr bool is_max = std::is_same_v; + constexpr bool is_sum = std::is_same_v; + if constexpr (std::is_same_v && (is_max || is_sum)) { + if constexpr (ROWS == 16) { + sk_float16v v; +#pragma unroll + for (int r = 0; r < 16; r++) v[r] = a[r]; + return is_max ? sk_hreduce16_max(v) : sk_hreduce16_add(v); + } else { + sk_float8v v; +#pragma unroll + for (int r = 0; r < 8; r++) v[r] = a[r]; + return is_max ? sk_hreduce8_max(v) : sk_hreduce8_add(v); + } + } +#endif +#pragma unroll + for (int s = ROWS / 2; s >= 1; s >>= 1) { + const bool hi = (lane & s) != 0; +#pragma unroll + for (int r = 0; r < s; r++) { + const type keep = hi ? a[r + s] : a[r]; + const type send = hi ? a[r] : a[r + s]; + a[r] = op::template op(keep, packed_shfl_sync(MASK_ALL, send, lane ^ s)); + } + } + if constexpr (ROWS == 8) // rows 0-7 are spread over both lane halves: fold them + a[0] = op::template op(a[0], packed_shfl_sync(MASK_ALL, a[0], lane ^ 8)); + return a[0]; +} + template static inline void row_reduce(V &row_accum, const T &src, const V &src_accum) { // I actually like these static asserts because they give more verbose errors @@ -66,56 +214,38 @@ static inline void row_reduce(V &row_accum, const T &src, const V &src_accum) { for (int i = 0; i < src.height; i++) { #pragma unroll for (int j = 0; j < src.width; j++) { - type tile_result{}; + // Per-lane row partials: rows[2k] = data[k].x, rows[2k+1] = data[k].y + // (the 32-column shapes first combine their two 16-column halves in-lane), + // then one transposing butterfly leaves row L's reduction in lane L. + type tile_result; if constexpr (std::is_same_v) { - // 8x32: 8 rows, data[0-3]=cols0-15, data[4-7]=cols16-31 (offset +4, - // not +8). Only 4 packed pairs -> fills lanes 0-7 (rows 0-7). + // 8x32: 8 rows, data[0-3]=cols0-15, data[4-7]=cols16-31 (offset +4). + type rows[8]; #pragma unroll for (int k = 0; k < 4; k++) { - type val_x = op::template op(src.tiles[i][j].data[k].x(), - src.tiles[i][j].data[k + 4].x()); - type val_y = op::template op(src.tiles[i][j].data[k].y(), - src.tiles[i][j].data[k + 4].y()); - // XOR butterfly reduction across 16 SIMD lanes (32 cols -> scalar) -#pragma unroll - for (int stride = 8; stride >= 1; stride >>= 1) { - val_x = op::template op( - val_x, packed_shfl_sync(MASK_ALL, val_x, lane ^ stride)); - val_y = op::template op( - val_y, packed_shfl_sync(MASK_ALL, val_y, lane ^ stride)); - } - if (lane == 2 * k) tile_result = val_x; - if (lane == 2 * k + 1) tile_result = val_y; + rows[2 * k] = op::template op(src.tiles[i][j].data[k].x(), + src.tiles[i][j].data[k + 4].x()); + rows[2 * k + 1] = op::template op(src.tiles[i][j].data[k].y(), + src.tiles[i][j].data[k + 4].y()); } + tile_result = transpose_reduce_rows(rows, lane); } else { + type rows[16]; #pragma unroll - for (int k = 0; k < 8; k++) { - type val_x, val_y; - if constexpr (std::is_same_v) { - // 16x32: combine two 16-column halves within-lane first - val_x = op::template op(src.tiles[i][j].data[k].x(), - src.tiles[i][j].data[k + 8].x()); - val_y = op::template op(src.tiles[i][j].data[k].y(), - src.tiles[i][j].data[k + 8].y()); - } else { - val_x = src.tiles[i][j].data[k].x(); - val_y = src.tiles[i][j].data[k].y(); - } - // XOR butterfly reduction across 16 SIMD lanes -#pragma unroll - for (int stride = 8; stride >= 1; stride >>= 1) { - val_x = op::template op( - val_x, packed_shfl_sync(MASK_ALL, val_x, lane ^ stride)); - val_y = op::template op( - val_y, packed_shfl_sync(MASK_ALL, val_y, lane ^ stride)); + for (int k = 0; k < 8; k++) { + if constexpr (std::is_same_v) { + rows[2 * k] = op::template op(src.tiles[i][j].data[k].x(), + src.tiles[i][j].data[k + 8].x()); + rows[2 * k + 1] = op::template op(src.tiles[i][j].data[k].y(), + src.tiles[i][j].data[k + 8].y()); + } else { + rows[2 * k] = src.tiles[i][j].data[k].x(); + rows[2 * k + 1] = src.tiles[i][j].data[k].y(); + } } - // All lanes now have reduce(row 2k) in val_x, reduce(row 2k+1) in val_y. - // Scatter: lane 2k stores reduce(row 2k), lane 2k+1 stores reduce(row 2k+1). - if (lane == 2 * k) tile_result = val_x; - if (lane == 2 * k + 1) tile_result = val_y; - } + tile_result = transpose_reduce_rows(rows, lane); } if (j == 0) { if constexpr (reset) { diff --git a/kernels/attention/gqa_forward.dp.cpp b/kernels/attention/gqa_forward.dp.cpp index 9592613..a9942db 100644 --- a/kernels/attention/gqa_forward.dp.cpp +++ b/kernels/attention/gqa_forward.dp.cpp @@ -36,6 +36,8 @@ #include "kittens.dp.hpp" #include "ops/warp/register/tile/conversions.dp.hpp" +#include "ops/warp/memory/tile/gl_desc.dp.hpp" +#include using namespace kittens; @@ -111,6 +113,14 @@ struct fwd_cfg { // load->MMA fences alone cost up to ~15% at large-N causal. Do NOT re-add them. // ============================================================ +// Calls f(std::integral_constant) for I = 0..N-1, so the block-load +// offsets below stay compile-time immediates. +template inline void static_for(F &&f) { + [&](std::integer_sequence) { + (f(std::integral_constant{}), ...); + }(std::make_integer_sequence{}); +} + // ============================================================ // Kernel (v16 compute core, parametrised on WG / SEQ_Q / CAUSAL) // ============================================================ @@ -176,16 +186,53 @@ SYCL_EXTERNAL void attend_ker_gqa_fwd( loop_kv_tiles = num_kv_tiles; } - // ── Prefetch first K tile ────────────────────────────────── - k_chunk_t k_pf; - #pragma unroll - for (int idx = 0; idx < PF && idx < loop_kv_tiles; idx++) { - #pragma unroll - for (int c = 0; c < inner_loop; c++) { - kittens::prefetch_load_part<1>(k_pf, Kg, - {batch_id, idx, kv_head_id, 0}, {0, c}); + // ── K/V block-load payloads, built once (sycl-tla style) ── + // Each (batch, kv_head) slice of K/V is an N_kv x D surface; inside the loop + // only the row (Y) is updated, every other offset is an immediate. Shapes + // match the flat loads they replace: K^T as d32 8x16 transposed blocks, V as + // VNNI 16x32x2, prefetches as 32x32 bf16. + using kt_desc_t = gl_desc<32, 16, 8, 1, false, true>; + using v_desc_t = gl_desc<16, 32, 16, 2, true>; + using pf_desc_t = gl_desc<16, 32, 32, 1>; + constexpr int v_cols = D / (SPLIT_D ? 2 : 1); + // The slice offset and pitch come from each gl's own stride accessor -- the + // same one global_to_register and friends use -- so this site cannot drift + // from the layout the rest of the library assumes. K and V are derived + // separately: sharing one set of numbers was only correct because the two + // happen to be constructed with identical dims. + const long k_off = (long)batch_id * (long)Kg.template stride<0>() + + (long)kv_head_id * (long)Kg.template stride<2>(); + const long v_off = (long)batch_id * (long)Vg.template stride<0>() + + (long)kv_head_id * (long)Vg.template stride<2>(); + const int k_pitch = (int)(Kg.template stride<1>() * sizeof(bf16)); + const int v_pitch = (int)(Vg.template stride<1>() * sizeof(bf16)); + const int k_rows = Kg.depth(); + const int v_rows = Vg.depth(); + kt_desc_t k_desc; pf_desc_t k_pf; + v_desc_t v_desc; pf_desc_t v_pf; + k_desc.init(Kg.raw_ptr + k_off, D * (int)sizeof(bf16), k_rows, k_pitch); + k_pf.init (Kg.raw_ptr + k_off, D * (int)sizeof(bf16), k_rows, k_pitch); + v_desc.init(Vg.raw_ptr + v_off, D * (int)sizeof(bf16), v_rows, v_pitch); + v_pf.init (Vg.raw_ptr + v_off, D * (int)sizeof(bf16), v_rows, v_pitch); + + // Cooperative prefetch: every subgroup of the WG reads the same K/V tile, + // so the next tile's pf_blocks 32x32 blocks (K then V columns) are split + // across the subgroups instead of each subgroup prefetching all of them. + constexpr int pf_blocks = 2 * (D / kv_stride); + auto prefetch_tile = [&](int tile) { + for (int b = sg_id % pf_blocks; b < pf_blocks; b += num_sg) { + if (b < pf_blocks / 2) { + k_pf.set_block(b * kv_stride, tile * kv_stride); + k_pf.prefetch(); + } else { + v_pf.set_block((b - pf_blocks / 2) * kv_stride, tile * kv_stride); + v_pf.prefetch(); + } } - } + }; + + // ── Prefetch first K/V tile ──────────────────────────────── + if (PF > 0 && loop_kv_tiles > 0) prefetch_tile(0); int prefetch = PF; @@ -199,17 +246,21 @@ SYCL_EXTERNAL void attend_ker_gqa_fwd( zero(att_block); max_vec_last = max_vec; + k_desc.set_y(i * kv_stride); + v_desc.set_y(i * kv_stride); + if (prefetch < loop_kv_tiles) prefetch_tile(prefetch); + // ── Step 1: S[seq_q×32] = Q[seq_q×D] × K^T[D×32] ────── k_chunk_t k_reg; - #pragma unroll - for (int c = 0; c < inner_loop; c++) { - kittens::load_transpose_part<1>(k_reg, Kg, - {batch_id, i, kv_head_id, 0}, {0, c}); - if (prefetch < loop_kv_tiles) - kittens::prefetch_load_part<1>(k_reg, Kg, - {batch_id, prefetch, kv_head_id, 0}, {0, c}); + static_for([&](auto c) { + constexpr int x = decltype(c)::value * kv_stride / 2; // d32 column + int *kd = reinterpret_cast(&k_reg.tiles[0][0].data[0]); + k_desc.template load(kd); + k_desc.template load(kd + 8); + k_desc.template load(kd + 16); + k_desc.template load(kd + 24); mma(att_block, q_preloaded[c], k_reg, att_block); - } + }); // ── Step 2: Causal mask ──────────────────────────────── if constexpr (CAUSAL_T) { @@ -242,25 +293,19 @@ SYCL_EXTERNAL void attend_ker_gqa_fwd( p_mma_t p_mma; p_mma = att_block; // float→bf16 conversion + // V columns [h*v_cols, (h+1)*v_cols) into v_reg. v_half_t v_reg; + auto load_v = [&](auto h) { + static_for([&](auto j) { + v_desc.template load( + &v_reg.tiles[0][decltype(j)::value].data[0]); + }); + }; + load_v(std::integral_constant{}); + mma(o_reg_1, p_mma, v_reg, o_reg_1); if constexpr (SPLIT_D) { - kittens::load_part<1>(v_reg, Vg, {batch_id, i, kv_head_id, 0}, {0, 0}); - mma(o_reg_1, p_mma, v_reg, o_reg_1); - if (prefetch < loop_kv_tiles) - kittens::prefetch_load_part<1>(v_reg, Vg, - {batch_id, prefetch, kv_head_id, 0}, {0, 0}); - kittens::load_part<1>(v_reg, Vg, - {batch_id, i, kv_head_id, 0}, {0, second_index_start}); + load_v(std::integral_constant{}); mma(o_reg_2, p_mma, v_reg, o_reg_2); - if (prefetch < loop_kv_tiles) - kittens::prefetch_load_part<1>(v_reg, Vg, - {batch_id, prefetch, kv_head_id, 0}, - {0, second_index_start}); - } else { - kittens::load<1>(v_reg, Vg, {batch_id, i, kv_head_id, 0}); - mma(o_reg_1, p_mma, v_reg, o_reg_1); - if (prefetch < loop_kv_tiles) - kittens::prefetch_load<1>(v_reg, Vg, {batch_id, prefetch, kv_head_id, 0}); } } @@ -355,6 +400,12 @@ static inline fwd_choice choose_fwd(int N, int B, bool causal, bool bf16_out) { // default. bf16 and f32 winners diverge on 22 shapes (the bf16 store // epilogue shifts the register-pressure balance), so we keep two tables. // N>2048 is uniformly wg256/sq16 (both dtypes, both causal, all B). +#ifdef KITTENS_BMG + // B70 (Xe2) sweep, bf16 out, B in {1,4}, N 512..8192: with the register- + // region row reductions wg256/sq8 is fastest (or within noise) for every + // causal and non-causal shape (non-causal N4096: 77 vs 50 TF for wg256/sq16). + if (bf16_out && N >= 512) return fwd_choice{256, 8}; +#endif if (N > 2048) return fwd_choice{256, 16}; if (bf16_out) { if (causal) { diff --git a/kernels/common.mk b/kernels/common.mk index 1b27f76..9abf657 100644 --- a/kernels/common.mk +++ b/kernels/common.mk @@ -6,7 +6,14 @@ CXX := icpx endif PYTHON ?= python3 -GPU_TARGET ?= intel_gpu_pvc +# AOT target. When not given, detect the local Intel GPU (see detect_gpu_target.sh); +# fall back to PVC when none is found, e.g. on a GPU-less build host. +ifeq ($(origin GPU_TARGET), undefined) +GPU_TARGET := $(shell sh $(SK_ROOT)/kernels/detect_gpu_target.sh "$(CXX)") +ifeq ($(GPU_TARGET),) +GPU_TARGET := intel_gpu_pvc +endif +endif BUILD_DIR ?= build SYCL_TARGET_FLAGS ?= -fsycl-targets=$(GPU_TARGET) @@ -18,8 +25,22 @@ ifeq ($(filter $(CONFIG),pytorch standalone),) $(error CONFIG must be pytorch or standalone, got '$(CONFIG)') endif +# 256-GRF ("large register file") AOT mode. PVC builds pick 256 GRF automatically +# (-autoGRFSelection); Battlemage (Xe2) AOT defaults to 128 GRF and the PVC-sized +# register tiles then spill silently (B70: GEMM 4096^3 4.4 -> 148 TF, GQA forward +# 3.6x). LARGE_GRF=auto (default) enables it for bmg targets; 1/0 force on/off. +LARGE_GRF ?= auto +ifeq ($(LARGE_GRF),auto) +SK_LARGE_GRF := $(if $(findstring bmg,$(GPU_TARGET)),1,0) +else +SK_LARGE_GRF := $(LARGE_GRF) +endif +LARGE_GRF_FLAGS := $(if $(filter 1,$(SK_LARGE_GRF)),-Xs "-options -ze-opt-large-register-file") +# Lets kernels pick Battlemage-specific launch configs (e.g. gqa_forward choose_fwd). +TARGET_DEFS := $(if $(findstring bmg,$(GPU_TARGET)),-DKITTENS_BMG) + SPIRV_FLAGS := -Xspirv-translator -spirv-ext=+SPV_INTEL_split_barrier,+SPV_INTEL_2d_block_io,+SPV_INTEL_subgroup_matrix_multiply_accumulate -COMMON_SYCL_FLAGS := -fsycl $(SYCL_TARGET_FLAGS) -std=c++20 -O3 -fsycl-unnamed-lambda -Wno-deprecated-declarations $(SPIRV_FLAGS) -DKITTENS_INTEL -DKITTENS_XE -DKITTENS_FAST_BF16_CONVERT -I$(SK_ROOT)/include +COMMON_SYCL_FLAGS := -fsycl $(SYCL_TARGET_FLAGS) $(LARGE_GRF_FLAGS) $(TARGET_DEFS) -std=c++20 -O3 -fsycl-unnamed-lambda -Wno-deprecated-declarations $(SPIRV_FLAGS) -DKITTENS_INTEL -DKITTENS_XE -DKITTENS_FAST_BF16_CONVERT -I$(SK_ROOT)/include PYTHON_INCLUDE = $(shell "$(PYTHON)" -c "import sysconfig; print(sysconfig.get_path('include'))") PYBIND_INCLUDE = $(shell "$(PYTHON)" -c "import pybind11; print(pybind11.get_include())") @@ -32,6 +53,7 @@ EXTENSION_FLAGS = $(COMMON_SYCL_FLAGS) -shared -fPIC -DTORCH_COMPILE -DTORCH_API check-sycl: @command -v "$(CXX)" >/dev/null 2>&1 || { echo "ERROR: icpx not found. Source env.src first."; exit 1; } + @echo "SyclKittens: GPU_TARGET=$(GPU_TARGET) LARGE_GRF=$(SK_LARGE_GRF)" check-python: check-sycl @"$(PYTHON)" -c "import pathlib, pybind11, torch; from torch.utils.cpp_extension import include_paths; assert any((pathlib.Path(p) / 'c10/xpu/XPUStream.h').is_file() for p in include_paths()), 'XPU-enabled PyTorch headers are required'" diff --git a/kernels/decode/Makefile b/kernels/decode/Makefile index 7c33632..70a1741 100644 --- a/kernels/decode/Makefile +++ b/kernels/decode/Makefile @@ -1,3 +1,6 @@ +# Decode is bandwidth-bound: 256 GRF halves the threads per XVE and costs 7-11% +# per token on B70, so keep 128 GRF unless LARGE_GRF=1 is passed explicitly. +LARGE_GRF ?= 0 include ../common.mk .PHONY: all decode diff --git a/kernels/detect_gpu_target.sh b/kernels/detect_gpu_target.sh new file mode 100644 index 0000000..42ff13b --- /dev/null +++ b/kernels/detect_gpu_target.sh @@ -0,0 +1,23 @@ +#!/bin/sh +# Print the AOT target (-fsycl-targets value) of the first local Intel GPU, or +# nothing if none is found. Used by common.mk when GPU_TARGET is not set. +# usage: detect_gpu_target.sh [compiler] +# Tries the compiler's own sycl-ls first ("Architecture: intel_gpu_*"), then the +# PCI device id in sysfs, which works without a GPU runtime. + +cxx=$(command -v "${1:-icpx}" 2>/dev/null) +for sl in "${cxx:+$(dirname "$cxx")/sycl-ls}" "$(command -v sycl-ls 2>/dev/null)"; do + [ -n "$sl" ] && [ -x "$sl" ] || continue + arch=$("$sl" --verbose 2>/dev/null | + sed -n 's/^ *Architecture: *\(intel_gpu_[a-z0-9_]*\).*/\1/p' | head -n 1) + [ -n "$arch" ] && { echo "$arch"; exit 0; } +done + +for d in /sys/class/drm/card*/device; do + [ "$(cat "$d/vendor" 2>/dev/null)" = 0x8086 ] || continue + case "$(cat "$d/device" 2>/dev/null)" in + 0x0bd*) echo intel_gpu_pvc; exit 0 ;; # Data Center GPU Max (PVC) + 0xe20*|0xe21*) echo intel_gpu_bmg_g21; exit 0 ;; # Arc B580/B570 (BMG-G21) + 0xe22*) echo intel_gpu_bmg_g31; exit 0 ;; # Arc Pro B70/B60 (BMG-G31) + esac +done