From 921063cd7ce30a9f12a87465299ca1112b0a13bc Mon Sep 17 00:00:00 2001 From: "Chen, Sheng S" Date: Mon, 12 Oct 2026 23:17:30 -0700 Subject: [PATCH 1/8] build: detect the local GPU target, 256-GRF mode on Battlemage When GPU_TARGET is not given, common.mk detects the local Intel GPU (detect_gpu_target.sh: sycl-ls, then the sysfs PCI id; PVC fallback). LARGE_GRF=auto adds -ze-opt-large-register-file on bmg targets: Xe2 AOT defaults to 128 GRF and the PVC-sized tiles spill there. bmg builds also get -DKITTENS_BMG for Battlemage launch configs. decode keeps 128 GRF (LARGE_GRF ?= 0), it is bandwidth-bound. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Chen, Sheng S --- kernels/common.mk | 26 ++++++++++++++++++++++++-- kernels/decode/Makefile | 3 +++ kernels/detect_gpu_target.sh | 23 +++++++++++++++++++++++ 3 files changed, 50 insertions(+), 2 deletions(-) create mode 100644 kernels/detect_gpu_target.sh 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 From f94e4b85e306b532663f023e1c7b5ac922d91162 Mon Sep 17 00:00:00 2001 From: "Chen, Sheng S" Date: Mon, 12 Oct 2026 23:17:30 -0700 Subject: [PATCH 2/8] base_types: bf16_2 conversion reuses the scalar converter The bf16_2 vec constructor lowers to a software round-to-nearest-even sequence (shr/and/add3 plus a NaN branch) on Xe2; converting the two lanes with the scalar converter emits the hardware conversion. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Chen, Sheng S --- include/common/base_types.dp.hpp | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) 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 { From 2550b229971b3989a68800c37121f46226dbfd59 Mon Sep 17 00:00:00 2001 From: "Chen, Sheng S" Date: Mon, 12 Oct 2026 23:17:30 -0700 Subject: [PATCH 3/8] reductions: transposing butterfly for register-tile row reductions row_reduce now reduces all rows of a base tile with one transposing butterfly (lane L ends up with row L): 15 shuffles for 16 rows instead of 64. Shuffles lower to serialized a0-indirect moves on Xe, so this was the main stall in the attention softmax. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Chen, Sheng S --- .../ops/warp/register/tile/reductions.dp.hpp | 93 ++++++++++--------- 1 file changed, 51 insertions(+), 42 deletions(-) diff --git a/include/ops/warp/register/tile/reductions.dp.hpp b/include/ops/warp/register/tile/reductions.dp.hpp index 5e6f01f..fc7e6ac 100644 --- a/include/ops/warp/register/tile/reductions.dp.hpp +++ b/include/ops/warp/register/tile/reductions.dp.hpp @@ -37,6 +37,33 @@ 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. */ +/** + * @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"); +#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 +93,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) { From 4fdd0b8c620372472b6656632e714d81d7d64b69 Mon Sep 17 00:00:00 2001 From: "Chen, Sheng S" Date: Mon, 12 Oct 2026 23:17:30 -0700 Subject: [PATCH 4/8] reductions: register-region transpose-reduce for float max/sum On device, float max/sum row reductions of 16- and 8-row tiles use inline vISA register regions (the data movement of sycl-tla's hreduce16/8, BSD-3-Clause) instead of shuffles, so the reduction needs no a0-indirect moves at all. The shuffle butterfly stays as the generic fallback. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Chen, Sheng S --- .../ops/warp/register/tile/reductions.dp.hpp | 121 ++++++++++++++++++ 1 file changed, 121 insertions(+) diff --git a/include/ops/warp/register/tile/reductions.dp.hpp b/include/ops/warp/register/tile/reductions.dp.hpp index fc7e6ac..4d5ed32 100644 --- a/include/ops/warp/register/tile/reductions.dp.hpp +++ b/include/ops/warp/register/tile/reductions.dp.hpp @@ -37,6 +37,110 @@ 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). @@ -49,6 +153,23 @@ inline float reduce_16x16(const sycl::float2 *first_half, 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; From b2cf57e2803c5aaa42569815255882637c30a17e Mon Sep 17 00:00:00 2001 From: "Chen, Sheng S" Date: Mon, 12 Oct 2026 23:17:30 -0700 Subject: [PATCH 5/8] gqa_forward: use wg256/seq_q8 on Battlemage A WG x seq_q sweep on B70 (bf16 out, B in {1,4}, N 512..8192) with the register-region reductions shows wg256/seq_q8 is fastest, or within noise, for every causal and non-causal shape. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Chen, Sheng S --- kernels/attention/gqa_forward.dp.cpp | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/kernels/attention/gqa_forward.dp.cpp b/kernels/attention/gqa_forward.dp.cpp index 9592613..7c971b9 100644 --- a/kernels/attention/gqa_forward.dp.cpp +++ b/kernels/attention/gqa_forward.dp.cpp @@ -355,6 +355,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) { From 5e921a4a7399ea8adc64e0ad6d4bdfe05ddd1fb7 Mon Sep 17 00:00:00 2001 From: "Chen, Sheng S" Date: Mon, 12 Oct 2026 23:17:30 -0700 Subject: [PATCH 6/8] B70 attention: hoisted 2D block-load payloads for K/V The forward KV loop now builds one block2d address payload per K/V stream (K^T, K prefetch, V, V prefetch) before the loop and only updates the row per KV tile; all other offsets are vISA immediates. This removes the per-load payload rebuild (364 -> 4 scalar movs per 64 DPAS; loop 504 -> 280 instrs). gl_desc gains a transposed mode, a raw-surface init, compile-time load and prefetch offsets, and int storage for VNNI reads (short storage made IGC re-interleave every word). B70, H32/8 D128 bf16: causal N4096 83 -> 89.5 TF, N8192 71.5 -> 86 TF, non-causal N4096 77 -> 90 TF; all 11 standalone harnesses pass. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Chen, Sheng S --- include/ops/warp/memory/tile/gl_desc.dp.hpp | 63 +++++++----- kernels/attention/gqa_forward.dp.cpp | 102 +++++++++++++------- 2 files changed, 110 insertions(+), 55 deletions(-) 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/kernels/attention/gqa_forward.dp.cpp b/kernels/attention/gqa_forward.dp.cpp index 7c971b9..0d3a9fd 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,30 @@ SYCL_EXTERNAL void attend_ker_gqa_fwd( loop_kv_tiles = num_kv_tiles; } + // ── 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); + const long kv_off = ((long)batch_id * Kg.depth() * Kg.rows() + kv_head_id) * D; + const int kv_pitch = Kg.rows() * D * (int)sizeof(bf16); + const int kv_rows = Kg.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 + kv_off, D * (int)sizeof(bf16), kv_rows, kv_pitch); + k_pf.init (Kg.raw_ptr + kv_off, D * (int)sizeof(bf16), kv_rows, kv_pitch); + v_desc.init(Vg.raw_ptr + kv_off, D * (int)sizeof(bf16), kv_rows, kv_pitch); + v_pf.init (Vg.raw_ptr + kv_off, D * (int)sizeof(bf16), kv_rows, kv_pitch); + // ── 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}); - } - } + if (PF > 0 && loop_kv_tiles > 0) + static_for([&](auto c) { + k_pf.template prefetch(); + }); int prefetch = PF; @@ -199,17 +223,26 @@ SYCL_EXTERNAL void attend_ker_gqa_fwd( zero(att_block); max_vec_last = max_vec; + const bool do_pf = prefetch < loop_kv_tiles; + k_desc.set_y(i * kv_stride); + v_desc.set_y(i * kv_stride); + if (do_pf) { + k_pf.set_y(prefetch * kv_stride); + v_pf.set_y(prefetch * kv_stride); + } + // ── 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); + if (do_pf) k_pf.template prefetch(); mma(att_block, q_preloaded[c], k_reg, att_block); - } + }); // ── Step 2: Causal mask ──────────────────────────────── if constexpr (CAUSAL_T) { @@ -242,25 +275,28 @@ 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, then prefetch the same + // columns of the next tile. 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]); + }); + }; + auto prefetch_v = [&](auto h) { + if (do_pf) + static_for([&](auto j) { + v_pf.template prefetch(); + }); + }; + load_v(std::integral_constant{}); + mma(o_reg_1, p_mma, v_reg, o_reg_1); + prefetch_v(std::integral_constant{}); 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}); + prefetch_v(std::integral_constant{}); } } From 27b206a414c3924da8bb05fdecd625c7d3935dc5 Mon Sep 17 00:00:00 2001 From: "Chen, Sheng S" Date: Mon, 12 Oct 2026 23:17:30 -0700 Subject: [PATCH 7/8] B70 attention: cooperative K/V prefetch across subgroups Every subgroup of a forward WG reads the same K/V tile, yet each one prefetched the whole next tile (8 32x32 blocks per KV tile). The next tile's blocks are now split across the subgroups, one block each for WG256, which cuts the per-subgroup 2D sends per KV tile from 28 to 21. B70, H32/8 D128 bf16: causal N4096 88 -> 96 TF, N8192 86 -> 91 TF, B4 N2048 causal 78 -> 82 TF / full 83 -> 87.5 TF, N1024 +3-4%; non-causal N4096 flat (90 -> 89). A 64-key iteration (half the reductions/rescales), early V loads and a lazy O rescale were tried and dropped: they were slower or neutral, the loop is load-latency bound rather than ALU bound. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Chen, Sheng S --- kernels/attention/gqa_forward.dp.cpp | 41 ++++++++++++++-------------- 1 file changed, 20 insertions(+), 21 deletions(-) diff --git a/kernels/attention/gqa_forward.dp.cpp b/kernels/attention/gqa_forward.dp.cpp index 0d3a9fd..f8d9e58 100644 --- a/kernels/attention/gqa_forward.dp.cpp +++ b/kernels/attention/gqa_forward.dp.cpp @@ -205,11 +205,24 @@ SYCL_EXTERNAL void attend_ker_gqa_fwd( v_desc.init(Vg.raw_ptr + kv_off, D * (int)sizeof(bf16), kv_rows, kv_pitch); v_pf.init (Vg.raw_ptr + kv_off, D * (int)sizeof(bf16), kv_rows, kv_pitch); - // ── Prefetch first K tile ────────────────────────────────── - if (PF > 0 && loop_kv_tiles > 0) - static_for([&](auto c) { - k_pf.template prefetch(); - }); + // 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; @@ -223,13 +236,9 @@ SYCL_EXTERNAL void attend_ker_gqa_fwd( zero(att_block); max_vec_last = max_vec; - const bool do_pf = prefetch < loop_kv_tiles; k_desc.set_y(i * kv_stride); v_desc.set_y(i * kv_stride); - if (do_pf) { - k_pf.set_y(prefetch * kv_stride); - v_pf.set_y(prefetch * 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; @@ -240,7 +249,6 @@ SYCL_EXTERNAL void attend_ker_gqa_fwd( k_desc.template load(kd + 8); k_desc.template load(kd + 16); k_desc.template load(kd + 24); - if (do_pf) k_pf.template prefetch(); mma(att_block, q_preloaded[c], k_reg, att_block); }); @@ -275,8 +283,7 @@ 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, then prefetch the same - // columns of the next tile. + // 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) { @@ -284,19 +291,11 @@ SYCL_EXTERNAL void attend_ker_gqa_fwd( &v_reg.tiles[0][decltype(j)::value].data[0]); }); }; - auto prefetch_v = [&](auto h) { - if (do_pf) - static_for([&](auto j) { - v_pf.template prefetch(); - }); - }; load_v(std::integral_constant{}); mma(o_reg_1, p_mma, v_reg, o_reg_1); - prefetch_v(std::integral_constant{}); if constexpr (SPLIT_D) { load_v(std::integral_constant{}); mma(o_reg_2, p_mma, v_reg, o_reg_2); - prefetch_v(std::integral_constant{}); } } From c5c93baf8832d44bd6fef9c06f0203cf14095990 Mon Sep 17 00:00:00 2001 From: "Chen, Sheng S" Date: Mon, 19 Oct 2026 10:06:13 +0800 Subject: [PATCH 8/8] B70 attention: derive K/V block-load layout from gl strides The hoisted 2D block-load payloads open-coded the per-(batch, kv_head) slice arithmetic and then reused one set of numbers for both K and V, taking the offset, pitch and surface height from Kg alone. Derive each side from its own gl via the template stride<> accessor that global_to_register and the other memory ops already use, so this site cannot drift from the layout the rest of the library assumes and V no longer depends on Kg's dims. Numerically a no-op at the harness's dimensions, where Kg and Vg are constructed identically: residuals are unchanged (causal 1.56e-03, full 2.10e-03) and a 5-trial paired A/B shows causal +0.1% and full +0.5%, both inside the per-case noise floor. No new spills. Note that this does not make the path safe for padded or arbitrarily strided tensors, because gl cannot describe one: stride() is derived from the shape (gl.dp.hpp:171) rather than stored, and operator[] indexes densely the same way. Supporting real strides is a gl-level change, not a per-kernel one. Co-Authored-By: Claude Opus 5 --- kernels/attention/gqa_forward.dp.cpp | 24 +++++++++++++++++------- 1 file changed, 17 insertions(+), 7 deletions(-) diff --git a/kernels/attention/gqa_forward.dp.cpp b/kernels/attention/gqa_forward.dp.cpp index f8d9e58..a9942db 100644 --- a/kernels/attention/gqa_forward.dp.cpp +++ b/kernels/attention/gqa_forward.dp.cpp @@ -195,15 +195,25 @@ SYCL_EXTERNAL void attend_ker_gqa_fwd( 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); - const long kv_off = ((long)batch_id * Kg.depth() * Kg.rows() + kv_head_id) * D; - const int kv_pitch = Kg.rows() * D * (int)sizeof(bf16); - const int kv_rows = Kg.depth(); + // 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 + kv_off, D * (int)sizeof(bf16), kv_rows, kv_pitch); - k_pf.init (Kg.raw_ptr + kv_off, D * (int)sizeof(bf16), kv_rows, kv_pitch); - v_desc.init(Vg.raw_ptr + kv_off, D * (int)sizeof(bf16), kv_rows, kv_pitch); - v_pf.init (Vg.raw_ptr + kv_off, D * (int)sizeof(bf16), kv_rows, kv_pitch); + 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