From a9038c475ec7918a8da0d7be6482c48c3a791ff0 Mon Sep 17 00:00:00 2001 From: Yanwen Xu Date: Wed, 22 Jul 2026 12:58:25 -0700 Subject: [PATCH 01/28] Add ET_VK_EXECUTE_NODE_THRESHOLD opt-in GPU-watchdog workaround MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Large prefills (e.g. 8B @ 2048 tokens) pack >2.56s of GPU work into a single 128-node command buffer submission, tripping the sgpu job watchdog. Setting the env var submits every N nodes instead, without adding a blocking stall (submits are non-blocking; execute() fences once at the end) — a measurement aid until the real fix lands driver-side. --- .ci/scripts/build_llama_android.sh | 0 backends/vulkan/runtime/graph/ComputeGraph.cpp | 16 ++++++++++++++++ 2 files changed, 16 insertions(+) mode change 100644 => 100755 .ci/scripts/build_llama_android.sh diff --git a/.ci/scripts/build_llama_android.sh b/.ci/scripts/build_llama_android.sh old mode 100644 new mode 100755 diff --git a/backends/vulkan/runtime/graph/ComputeGraph.cpp b/backends/vulkan/runtime/graph/ComputeGraph.cpp index 3accdf375cb..6b9b292989f 100644 --- a/backends/vulkan/runtime/graph/ComputeGraph.cpp +++ b/backends/vulkan/runtime/graph/ComputeGraph.cpp @@ -176,6 +176,22 @@ ComputeGraph::ComputeGraph(GraphConfig config) config_.execute_initial_threshold_node_count = 64; } + // Opt-in GPU-watchdog workaround (default behaviour unchanged). Set + // ET_VK_EXECUTE_NODE_THRESHOLD=N to submit a new command buffer every N graph + // nodes instead of the default 128. A large prefill (e.g. 8B @ 2048 tokens) + // packs > 2.56 s of GPU work into a single 128-node submission, tripping the + // sgpu job watchdog (hard reset, lost run). Submits here are non-blocking + // (execute() defers and fences once at the end), so a smaller N only adds a + // little submit overhead, not a per-batch stall. TEMPORARY measurement aid — + // the real fix is driver-side. + if (const char* thr = std::getenv("ET_VK_EXECUTE_NODE_THRESHOLD")) { + const int n = std::atoi(thr); + if (n > 0) { + config_.execute_threshold_node_count = static_cast(n); + config_.execute_initial_threshold_node_count = static_cast(n); + } + } + // Check if the underlying GPU can access accelerated integer dot product // instructions can_use_int8_dot_product_ = From 745a7381bc1ae3756179308721ac838dcab184fb Mon Sep 17 00:00:00 2001 From: Sicheng Stephen Jia Date: Wed, 29 Jul 2026 16:08:56 -0400 Subject: [PATCH 02/28] Fix linear_dq8ca_q4gsw wrong output when M > K Differential Revision: D113962415 Pull Request resolved: https://github.com/pytorch/executorch/pull/21444 --- .../runtime/graph/ops/impl/QuantizedLinear.cpp | 12 ++++++++++-- .../vulkan/test/custom_ops/test_q4gsw_linear.cpp | 2 ++ 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp index b9a03d85161..45588e7e2e5 100644 --- a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp @@ -833,10 +833,18 @@ void quantized_linear_impl( num_groups = graph.size_at(-2, weight_scales_data); } + // Per-group int8 input sums buffer, indexed as ivec4[group_idx * M4 + m4] + // by both the producer (quantize_and_pack_4h4w_with_group_sums.glsl) and the + // consumer (linear_int8_input_sums_load.glslh). Capacity must therefore be + // num_groups * M4 ivec4 texels, sized by the input row count M -- NOT K. + // dtype is kInt to match the shaders' `int`/ivec4 binding (each texel is 4 + // int32 sums = 16 bytes). + const int64_t M = utils::val_at(-2, input_sizes); + const int64_t M4 = utils::div_up(M, int64_t(4)); TmpTensor int_input_sums( &graph, - {num_groups, K}, - graph.dtype_of(output), + {num_groups * M4 * 4}, + vkapi::kInt, utils::kBuffer, utils::kWidthPacked); diff --git a/backends/vulkan/test/custom_ops/test_q4gsw_linear.cpp b/backends/vulkan/test/custom_ops/test_q4gsw_linear.cpp index c1d66ab4aec..a9d01f0cbf0 100644 --- a/backends/vulkan/test/custom_ops/test_q4gsw_linear.cpp +++ b/backends/vulkan/test/custom_ops/test_q4gsw_linear.cpp @@ -269,6 +269,8 @@ std::vector generate_quantized_linear_test_cases() { {32, 64, 32, 16}, {32, 128, 64, 32}, {32, 256, 128, 64}, + {256, 128, 128, 32}, // M=256 > K=128; all dims < kRefDimSizeLimit + // Coopmat-eligible correctness shapes (M%64==0, N%64==0, K%32==0, // group_size%32==0). The Buffer+Half variant fires linear_q4gsw_coopmat / // linear_dq8ca_q4gsw_coopmat and is validated against the CPU reference. From f60080d04aaccb3796272cb2827c4a0922bb04e8 Mon Sep 17 00:00:00 2001 From: Yanwen Xu Date: Mon, 17 Aug 2026 18:37:38 -0700 Subject: [PATCH 03/28] [ET-VK] Port dev's shipped 4w/8da4w coopmat tiles + texture-IO WMMA to release/1.4 Aggregates the best-known-validated quantized-linear coopmat work, scattered across several worktree branches, onto a release/1.4 base: - Ports dev's e2e-validated shipped tile geometry for 4w (specs/036, 128x128x16) and 8da4w (specs/027, 64x32x32 dbuf2) verbatim into linear_qw_coopmat.glsl/yaml and linear_dq8ca_qw_coopmat.glsl/yaml. - Fixes two independent bugs found via ETDump (neither shader ever actually dispatched before these fixes, on any branch tested this session): - can_use_q4gsw_coopmat's dim>2 rank check rejected the real model's rank-3 [1, M, K] activations; replaced with dev's leading-dims-numel check (a [1, M, N] buffer is bit-identical to [M, N] when the leading dim is 1). - Q4gswLinear.cpp (upstream) silently hijacks the et_vk.linear_q4gsw.default op registration with its own tiled-only q4gsw_linear_gemm__* shaders, making QuantizedLinear.cpp's linear_q4gsw_coopmat path dead code for 4w. Same bug as memory quant-perf-rebase-orphaned-4w-coopmat; applied the same known fix (restore QuantizedLinear.cpp's registration, disable the colliding one). - Ports texture-dbuf4's texture-storage IO capability (specs/040/041) for both 4w and 8da4w, so WMMA can dispatch on the canonical texture3d _embq_ PTE instead of requiring the retired buffer-only export path. Gated behind ET_VK_TEXTURE_COOPMAT=1 + ET_VK_Q4GSW_COOPMAT_VARIANT=tsweep_dbuf4_t128x128k16g22s32 / ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4_t64x32k32g12s64 -- off by default, so behavior without the env vars is unchanged. ETDump-confirmed on M51 (secondary board, SUMD change-66942): 100% of 112 real prefill GEMM dispatches hit the coopmat kernel for both quant modes. Prefill tok/s (median of 5, Llama 3.2 1B, 2048-token prompt): 4w 1064.5 (1.43x vs release-1.4 tiled baseline), 8da4w 969.2 (1.27x). Correctness not yet run through the small-shape coopmat bench -- only "coherent, non-garbage output" observed so far. Authored with Claude. --- ...near_dq8ca_q4gsw_coopmat_tsweep_dbuf4.glsl | 634 +++++ ...near_dq8ca_q4gsw_coopmat_tsweep_dbuf4.yaml | 2164 +++++++++++++++++ .../ops/glsl/linear_dq8ca_qw_coopmat.glsl | 268 +- .../ops/glsl/linear_dq8ca_qw_coopmat.yaml | 18 +- .../linear_q4gsw_coopmat_tsweep_dbuf4.glsl | 473 ++++ .../linear_q4gsw_coopmat_tsweep_dbuf4.yaml | 858 +++++++ .../graph/ops/glsl/linear_qw_coopmat.yaml | 17 +- .../runtime/graph/ops/impl/Q4gswLinear.cpp | 8 +- .../graph/ops/impl/QuantizedLinear.cpp | 279 ++- backends/vulkan/runtime/vk_api/Adapter.h | 4 + 10 files changed, 4541 insertions(+), 182 deletions(-) create mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4.glsl create mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4.yaml create mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_coopmat_tsweep_dbuf4.glsl create mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_coopmat_tsweep_dbuf4.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4.glsl b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4.glsl new file mode 100644 index 00000000000..b8d4b309233 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4.glsl @@ -0,0 +1,634 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +/* + * TILE/SUBGROUP-SWEEP variant of the int8 dq8ca_q4gsw coopmat shader's dbuf4 + * ("store-first-for-next", the ORIGINAL loop structure before specs/025 User + * Story 1 picked dbuf2) loop structure (specs/041-dbuf4-tile-sweep). Forked + * from linear_dq8ca_q4gsw_coopmat_tsweep.glsl (which carries dbuf2's loop, + * the production winner) -- everything except the PROLOGUE/MAIN LOOP block + * is identical: bindings, spec-constants, tile-geometry templating, LDS + * layout (ColumnMajor B + skew), int8 WMMA thread maps, group epilog, + * bias/store epilogue. Only the loop structure is swapped to dbuf4, + * recovered from git commit 8d0f23ee78's + * linear_dq8ca_q4gsw_coopmat_dbuf4.glsl (see specs/041/reference/) -- the + * byte-identical pre-swap copy of what is now linear_dq8ca_qw_coopmat.glsl. + * + * The nested `groups x chunks` loop and unconditional group epilog are kept + * exactly as in dbuf2 -- flattening them crashes the Xclipse PAL compiler at + * large spec-resolved trip counts (see dbuf2's own header). Only the + * store/barrier/prefetch ORDER within each chunk iteration is inverted: + * + * dbuf2 (this file's base): store(temp, already prefetched -> cur slice) + * -> barrier -> MMA(cur) -> prefetch(next -> temp) [store owns the + * CURRENT chunk, at the iteration's start] + * dbuf4 (this file): barrier -> prefetch(next -> temp) -> MMA(cur) -> + * store(temp -> next slice) [store owns the NEXT chunk, at the + * iteration's end -- the mirror image] + * + * The group wsum/wsc ping-pong is inverted the same way: dbuf4 stores the + * next group's values (prefetched during the crossing chunk) at the TAIL of + * that chunk, instead of dbuf2's HEAD-of-new-group placement. + * + * Selected at dispatch via + * ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4_txkgs<32|64> + * (QuantizedLinear.cpp), additive to the existing tsweep_t... (dbuf2) + * namespace. + * + * KHR Cooperative Matrix variant of the dynamically-quantized-activation + * linear tiled shader (WEIGHT_NBITS=4): + * 4 -> linear_dq8ca_q4gsw_coopmat INT4 group-symmetric weight + * + * Performs: out[M,N] = dequant(int8_act) * dequant(int_w) (+ bias) + * via coopmat x coopmat -> coopmat on the matrix unit. + * + * Hard preconditions: + * M % WG_TILE_M == 0, N % WG_TILE_N == 0, K % WG_TILE_K == 0, + * INT4: group_size % WG_TILE_K == 0, + * device exposes coopmatx-> at 16x16x16. + */ + +#version 450 core + +#extension GL_KHR_cooperative_matrix : require +#extension GL_KHR_memory_scope_semantics : require +#extension GL_KHR_shader_subgroup_basic : enable +#extension GL_EXT_shader_explicit_arithmetic_types : require +#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_EXT_control_flow_attributes : enable + +#define PRECISION ${PRECISION} + +$if WEIGHT_NBITS == 4: + #define WEIGHT_INT4 + +$if HAS_BIAS: + #define HAS_BIAS + +$if WEIGHT_STORAGE == "buffer": + #define WEIGHT_BUFFER + +$if IO_STORAGE == "texture3d": + #define IO_TEXTURE + +layout(std430) buffer; + +#include "common.glslh" + +// Bindings — match add_linear_dqa_qw_node arg order: +// output(0), fp_input(1), packed_int8_input(2), int_input_sums(3 - unused), +// input_scales(4), input_zps(5), packed_weight(6), weight_sums(7), +// weight_scales(8), bias(9). +${layout_declare_tensor(B, "w", "t_output", "half", IO_STORAGE, is_scalar_array=True)} +// t_input is unread here -- the activations arrive already quantized in +// t_packed_int8_input -- but stays declared so the binding layout matches the +// dispatch site. It tracks IO_STORAGE so the two IO tensors stay consistent. +${layout_declare_tensor(B, "r", "t_input", "half", IO_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_packed_int8_input", "int", "buffer", is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_int8_input_sums", "int", "buffer", is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_int8_input_scales", "half", "texture3d")} +${layout_declare_tensor(B, "r", "t_int8_input_zps", "int8", "texture3d")} +${layout_declare_tensor(B, "r", "t_packed_weight", "int", WEIGHT_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_weight_sums", "int", "buffer", is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_weight_scales", "half", "buffer", is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_bias", "half", "buffer", is_scalar_array=True)} + +${layout_declare_ubo(B, "ivec4", "output_sizes")} +${layout_declare_ubo(B, "ivec4", "input_sizes")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +${layout_declare_spec_const(C, "int", "apply_bias", "0")} +// INT4 only; inert (0) for INT8 so the dispatcher's spec list lines up. +${layout_declare_spec_const(C, "int", "K4_per_group", "0")} +${layout_declare_spec_const(C, "int", "num_groups_arg", "0")} +${layout_declare_spec_const(C, "int", "out_N_arg", "0")} + +// Tile geometry +const uint MMA_M = ${MMA_M}; +const uint MMA_N = ${MMA_N}; +const uint MMA_K = ${MMA_K}; + +const uint WG_TILE_M = ${WG_TILE_M}; +const uint WG_TILE_N = ${WG_TILE_N}; +const uint WG_TILE_K = ${WG_TILE_K}; + +const uint SG_GRID_X = ${SG_GRID_X}; +const uint SG_GRID_Y = ${SG_GRID_Y}; +const uint SUBGROUP_SIZE = ${SUBGROUP_SIZE}; +const uint NUM_SUBGROUPS = SG_GRID_X * SG_GRID_Y; +const uint WG_SIZE = NUM_SUBGROUPS * SUBGROUP_SIZE; + +const uint SG_TILE_M = WG_TILE_M / SG_GRID_Y; +const uint SG_TILE_N = WG_TILE_N / SG_GRID_X; +const uint MMAS_PER_SG_M = SG_TILE_M / MMA_M; +const uint MMAS_PER_SG_N = SG_TILE_N / MMA_N; + +const uint A_SLAB_INT8 = WG_TILE_M * MMA_K; +const uint B_USEFUL_U32 = MMA_K / 4u; +const uint B_STRIDE_U32 = B_USEFUL_U32 + 1u; // +1 skew +const uint B_SLAB_U32 = WG_TILE_N * B_STRIDE_U32; +const uint NUM_K_SLABS = WG_TILE_K / MMA_K; + +const uint A_SLAB_U32 = A_SLAB_INT8 / 4u; +const uint A_STRIDE_U32 = MMA_K / 4u; + +// One ping-pong slice covers all K-slabs of one chunk. +const uint ASH_SLICE_U32 = NUM_K_SLABS * A_SLAB_U32; +const uint BSH_SLICE_U32 = NUM_K_SLABS * B_SLAB_U32; + +// Double-buffered MMA operand staging. +shared uint Ash_int8[2u * ASH_SLICE_U32]; +shared uint Bsh_int8[2u * BSH_SLICE_U32]; + +// Per-WG-tile-row activation params (loaded ONCE at WG start; constant +// across groups). +shared int izp_sh[WG_TILE_M]; // int32 (cast from int8 source) for broadcast +shared float ifs_sh[WG_TILE_M]; // float32 (cast from fp16 source) for broadcast + +// Per-(group, output-channel) weight params, ping-ponged by group parity. +// (For per-channel INT8 only slice 0 is ever used.) +shared int wsum_sh[2u * WG_TILE_N]; +shared float wsc_sh[2u * WG_TILE_N]; + +#ifdef HAS_BIAS +shared float bias_sh[WG_TILE_N]; +#endif + +#ifdef IO_TEXTURE +// Result staging for the imageStore epilogue, mirroring the fp16 kernel: +// SG_GRID_Y bands of MMA_M rows, each WG_TILE_N wide, row-major. A full +// WG_TILE_M x WG_TILE_N buffer would cost SG_GRID_Y/MMAS_PER_SG_M x more LDS +// and wreck occupancy. float16_t-typed because coopMatStore needs it. +const uint CSH_ROWS = SG_GRID_Y * MMA_M; +shared float16_t Csh_out[CSH_ROWS * WG_TILE_N]; +#endif + +// Running fp32 accumulator (across all groups). +coopmat + result[MMAS_PER_SG_M][MMAS_PER_SG_N]; + +// Per-group int32 MMA accumulator. +coopmat + accum_int32[MMAS_PER_SG_M][MMAS_PER_SG_N]; + +void main() { + const uvec2 tileID = uvec2(gl_WorkGroupID.xy); + const uvec2 warpInTile = uvec2( + gl_SubgroupID % SG_GRID_X, + gl_SubgroupID / SG_GRID_X); + + const uint K = uint(input_sizes.x); + const uint N = uint(output_sizes.x); + const uint N4 = (N + 3u) / 4u; + const uint nblocks_x_A = (K + 3u) >> 2u; + +#ifdef WEIGHT_INT4 + const uint num_groups = uint(num_groups_arg); + const uint CHUNKS_PER_GROUP = uint(K4_per_group) * 4u / WG_TILE_K; +#else + // Per-channel: a single quant "group" spanning all of K. The nested + // groups x chunks loop below collapses to a flat chunk loop, the wsum/wsc + // ping-pong never crosses a boundary, and the epilog runs exactly once. + const uint num_groups = 1u; + const uint CHUNKS_PER_GROUP = uint(num_groups_arg); +#endif + const uint num_chunks = num_groups * CHUNKS_PER_GROUP; + + const uint tile_m_start = WG_TILE_M * tileID.y; + const uint tile_n_start = WG_TILE_N * tileID.x; + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + result[i][j] = coopmat(0.0); + accum_int32[i][j] = coopmat(0); + } + } + + // --- A staging thread map: one (m4, k4) ivec4 block per active thread --- + const uint K_BLOCKS_PER_CHUNK = WG_TILE_K >> 2u; + const uint A_ACTIVE_THREADS = (WG_TILE_M >> 2u) * K_BLOCKS_PER_CHUNK; + const uint a_m_block = gl_LocalInvocationID.x / K_BLOCKS_PER_CHUNK; + const uint a_k_block = gl_LocalInvocationID.x % K_BLOCKS_PER_CHUNK; + const bool a_active = gl_LocalInvocationID.x < A_ACTIVE_THREADS; + +#ifdef WEIGHT_INT4 + // --- B staging thread map: (block, col) slots; each slot extracts one + // ColumnMajor LDS uint (4 K-contiguous sign-extended int8) --- + const uint B_TOTAL_SLOTS = K_BLOCKS_PER_CHUNK * WG_TILE_N; + const uint B_SLOTS_PER_THREAD = B_TOTAL_SLOTS / WG_SIZE; + const uint N8_PER_TILE = WG_TILE_N >> 3u; +#else + // --- B staging thread map: one (k4, n4) ivec4 block per active thread --- + const uint B_FETCH_SLOTS = K_BLOCKS_PER_CHUNK * (WG_TILE_N >> 2u); + const uint N4_PER_TILE = WG_TILE_N >> 2u; + const uint b_k4_in_chunk = gl_LocalInvocationID.x / N4_PER_TILE; + const uint b_n_uint_col = gl_LocalInvocationID.x % N4_PER_TILE; + const bool b_active = gl_LocalInvocationID.x < B_FETCH_SLOTS; +#endif + + // Prefetch temp registers. + ivec4 temp_A; +#ifdef WEIGHT_INT4 + ivec4 temp_B[B_SLOTS_PER_THREAD]; + int temp_wsum; + float temp_wsc; +#else + ivec4 temp_B; +#endif + + // ========================================================= + // PROLOGUE + // ========================================================= + if (gl_LocalInvocationID.x < (WG_TILE_M >> 2u)) { + const uint m4 = (tile_m_start >> 2u) + gl_LocalInvocationID.x; + const vec4 sc = vec4(texelFetch(t_int8_input_scales, ivec3(m4, 0, 0), 0)); + const ivec4 zp = texelFetch(t_int8_input_zps, ivec3(m4, 0, 0), 0); + const uint base = gl_LocalInvocationID.x * 4u; + ifs_sh[base + 0u] = sc.x; ifs_sh[base + 1u] = sc.y; + ifs_sh[base + 2u] = sc.z; ifs_sh[base + 3u] = sc.w; + izp_sh[base + 0u] = zp.x; izp_sh[base + 1u] = zp.y; + izp_sh[base + 2u] = zp.z; izp_sh[base + 3u] = zp.w; + } + // Group 0 weight sums/scales -> slice 0. + if (gl_LocalInvocationID.x < WG_TILE_N) { + const uint n_idx = tile_n_start + gl_LocalInvocationID.x; + f16vec4 sv = t_weight_scales[n_idx >> 2u]; + wsc_sh[gl_LocalInvocationID.x] = float(sv[n_idx & 3u]); + wsum_sh[gl_LocalInvocationID.x] = t_weight_sums[n_idx]; + } + memoryBarrierShared(); + barrier(); + + // izp/ifs are per-row activation params, constant across K groups — + // broadcast them into coopmats ONCE; the group epilog reuses them every + // group (they depend only on the row block i, not on the group or j). + coopmat + izp_bcast[MMAS_PER_SG_M]; + coopmat + ifs_bcast[MMAS_PER_SG_M]; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + const uint local_m_base = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + coopMatLoad( + izp_bcast[i], izp_sh, + local_m_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutColumnMajor); + coopMatLoad( + ifs_bcast[i], ifs_sh, + local_m_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutColumnMajor); + } + + // dbuf4: prefetch chunk 0 into temp registers, THEN store to slice 0 (no + // barrier here -- the main loop's first iteration barriers before + // reading slice 0). + if (a_active) { + const uint m4_global = (tile_m_start >> 2u) + a_m_block; + temp_A = t_packed_int8_input[m4_global * nblocks_x_A + a_k_block]; + } +#ifdef WEIGHT_INT4 + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { + const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; + const uint block_in_chunk = slot >> 3u; + const uint k4_blk = block_in_chunk / N8_PER_TILE; + const uint n8_blk = (tile_n_start >> 3u) + (block_in_chunk % N8_PER_TILE); +#ifdef WEIGHT_BUFFER + temp_B[si] = t_packed_weight[(n8_blk * nblocks_x_A) + k4_blk]; +#else + temp_B[si] = texelFetch(t_packed_weight, ivec2(k4_blk, n8_blk), 0); +#endif + } +#else + if (b_active) { + const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; +#ifdef WEIGHT_BUFFER + temp_B = t_packed_weight[(b_k4_in_chunk * N4) + block_x_w]; +#else + temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, b_k4_in_chunk), 0); +#endif + } +#endif + { + // store chunk 0 -> slice 0 + if (a_active) { + const uint slab_idx = a_k_block / (MMA_K >> 2u); + const uint k_uint_in_slab = a_k_block % (MMA_K >> 2u); + const uint base_row = a_m_block * 4u; + [[unroll]] for (uint m4i = 0; m4i < 4u; ++m4i) { + Ash_int8[slab_idx * A_SLAB_U32 + (base_row + m4i) * A_STRIDE_U32 + k_uint_in_slab] = + uint(temp_A[m4i]); + } + } +#ifdef WEIGHT_INT4 + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { + const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; + const uint block_in_chunk = slot >> 3u; + const uint col_in_block = slot & 7u; + const uint k4_in_chunk = block_in_chunk / N8_PER_TILE; + const uint n8_in_tile = block_in_chunk % N8_PER_TILE; + const uint r = col_in_block & 3u; + const uint parity = col_in_block >> 2u; + const int w = temp_B[si][r]; + const int base = int(4u * parity); + const int v0 = (((w >> (base + 0)) & 0xF) - 8) & 0xFF; + const int v1 = (((w >> (base + 8)) & 0xF) - 8) & 0xFF; + const int v2 = (((w >> (base + 16)) & 0xF) - 8) & 0xFF; + const int v3 = (((w >> (base + 24)) & 0xF) - 8) & 0xFF; + const uint n_col = n8_in_tile * 8u + r + parity * 4u; + const uint slab_idx = k4_in_chunk / (MMA_K >> 2u); + const uint k4_in_slab = k4_in_chunk % (MMA_K >> 2u); + Bsh_int8[slab_idx * B_SLAB_U32 + n_col * B_STRIDE_U32 + k4_in_slab] = + uint(v0 | (v1 << 8) | (v2 << 16) | (v3 << 24)); + } +#else + if (b_active) { + const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); + const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); + const uint n_col_base = b_n_uint_col * 4u; + [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { + Bsh_int8[slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab] = + uint(temp_B[n_in_blk]); + } + } +#endif + } + + // ========================================================= + // MAIN LOOP (dbuf4) — nested groups x chunks (kept nested; flattening it + // with a conditional coopmat epilog crashes the Xclipse PAL compiler at + // large spec-resolved trip counts). One barrier per chunk. Chunk + // iteration (global index `chunk`): + // 1. barrier — A/B slice (chunk%2) fully written; on the first chunk + // of group g, wsum/wsc slice (g%2) is too. + // 2. prefetch — chunk+1 (A blocks, B blocks) into temp; when chunk+1 + // starts a new group, also its wsum/wsc element. Skipped + // entirely on the final chunk. + // 3. int8 MMA — on slice (chunk%2) into accum_int32. + // 4. store — temp -> A/B slice ((chunk+1)%2), unpacking the weight; + // on a group boundary, wsum/wsc -> slice ((g+1)%2). + // The group epilog runs unconditionally at the tail of each group. + // ========================================================= + uint chunk = 0; + for (uint group_i = 0; group_i < num_groups; ++group_i) { + for (uint inner = 0; inner < CHUNKS_PER_GROUP; ++inner, ++chunk) { + const bool has_next = chunk + 1u < num_chunks; + const bool group_crossing = has_next && (inner + 1u == CHUNKS_PER_GROUP); + const uint cur_a = (chunk % 2u) * ASH_SLICE_U32; + const uint cur_b = (chunk % 2u) * BSH_SLICE_U32; + const uint nxt_a = ((chunk + 1u) % 2u) * ASH_SLICE_U32; + const uint nxt_b = ((chunk + 1u) % 2u) * BSH_SLICE_U32; + + barrier(); + + // --- 2. prefetch chunk+1 -> temp --- + if (has_next) { + const uint chunkK_nxt = (chunk + 1u) * WG_TILE_K; + if (a_active) { + const uint m4_global = (tile_m_start >> 2u) + a_m_block; + const uint k4_global = (chunkK_nxt >> 2u) + a_k_block; + temp_A = t_packed_int8_input[m4_global * nblocks_x_A + k4_global]; + } +#ifdef WEIGHT_INT4 + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { + const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; + const uint block_in_chunk = slot >> 3u; + const uint k4_blk = (chunkK_nxt >> 2u) + block_in_chunk / N8_PER_TILE; + const uint n8_blk = (tile_n_start >> 3u) + (block_in_chunk % N8_PER_TILE); +#ifdef WEIGHT_BUFFER + temp_B[si] = t_packed_weight[(n8_blk * nblocks_x_A) + k4_blk]; +#else + temp_B[si] = texelFetch(t_packed_weight, ivec2(k4_blk, n8_blk), 0); +#endif + } + if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { + const uint n_idx = tile_n_start + gl_LocalInvocationID.x; + f16vec4 sv = t_weight_scales[(group_i + 1u) * N4 + (n_idx >> 2u)]; + temp_wsc = float(sv[n_idx & 3u]); + temp_wsum = t_weight_sums[(group_i + 1u) * N + n_idx]; + } +#else + if (b_active) { + const uint block_y_w = (chunkK_nxt >> 2u) + b_k4_in_chunk; + const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; +#ifdef WEIGHT_BUFFER + temp_B = t_packed_weight[(block_y_w * N4) + block_x_w]; +#else + temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, block_y_w), 0); +#endif + } +#endif + } + + // --- 3. int8 MMA on the cur slice --- + [[unroll]] for (uint k = 0; k < NUM_K_SLABS; ++k) { + const uint slab_a_base_u32 = cur_a + k * A_SLAB_U32; + const uint slab_b_base_u32 = cur_b + k * B_SLAB_U32; + + coopmat matA[MMAS_PER_SG_M]; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + const uint row_a = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + coopMatLoad( + matA[i], Ash_int8, + slab_a_base_u32 + row_a * A_STRIDE_U32, + A_STRIDE_U32, + gl_CooperativeMatrixLayoutRowMajor); + } + + coopmat matB; + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint col_b = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopMatLoad( + matB, Bsh_int8, + slab_b_base_u32 + col_b * B_STRIDE_U32, + B_STRIDE_U32, + gl_CooperativeMatrixLayoutColumnMajor); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + accum_int32[i][j] = coopMatMulAdd(matA[i], matB, accum_int32[i][j]); + } + } + } + + // --- 4. store temp (chunk+1) -> nxt slice --- + if (has_next) { + if (a_active) { + const uint slab_idx = a_k_block / (MMA_K >> 2u); + const uint k_uint_in_slab = a_k_block % (MMA_K >> 2u); + const uint base_row = a_m_block * 4u; + [[unroll]] for (uint m4i = 0; m4i < 4u; ++m4i) { + Ash_int8[nxt_a + slab_idx * A_SLAB_U32 + (base_row + m4i) * A_STRIDE_U32 + k_uint_in_slab] = + uint(temp_A[m4i]); + } + } +#ifdef WEIGHT_INT4 + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { + const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; + const uint block_in_chunk = slot >> 3u; + const uint col_in_block = slot & 7u; + const uint k4_in_chunk = block_in_chunk / N8_PER_TILE; + const uint n8_in_tile = block_in_chunk % N8_PER_TILE; + const uint r = col_in_block & 3u; + const uint parity = col_in_block >> 2u; + const int w = temp_B[si][r]; + const int base = int(4u * parity); + const int v0 = (((w >> (base + 0)) & 0xF) - 8) & 0xFF; + const int v1 = (((w >> (base + 8)) & 0xF) - 8) & 0xFF; + const int v2 = (((w >> (base + 16)) & 0xF) - 8) & 0xFF; + const int v3 = (((w >> (base + 24)) & 0xF) - 8) & 0xFF; + const uint n_col = n8_in_tile * 8u + r + parity * 4u; + const uint slab_idx = k4_in_chunk / (MMA_K >> 2u); + const uint k4_in_slab = k4_in_chunk % (MMA_K >> 2u); + Bsh_int8[nxt_b + slab_idx * B_SLAB_U32 + n_col * B_STRIDE_U32 + k4_in_slab] = + uint(v0 | (v1 << 8) | (v2 << 16) | (v3 << 24)); + } + if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { + const uint wbase_nxt = ((group_i + 1u) % 2u) * WG_TILE_N; + wsum_sh[wbase_nxt + gl_LocalInvocationID.x] = temp_wsum; + wsc_sh[wbase_nxt + gl_LocalInvocationID.x] = temp_wsc; + } +#else + if (b_active) { + const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); + const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); + const uint n_col_base = b_n_uint_col * 4u; + [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { + Bsh_int8[nxt_b + slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab] = + uint(temp_B[n_in_blk]); + } + } +#endif + } + } // chunks + + // --- Group epilog: dequant accum_int32 -> result, reset accum --- + { + const uint wbase = (group_i % 2u) * WG_TILE_N; + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + + coopmat wsum_bcast; + coopMatLoad( + wsum_bcast, wsum_sh, + wbase + local_n_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutRowMajor); + + coopmat wsc_bcast; + coopMatLoad( + wsc_bcast, wsc_sh, + wbase + local_n_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutRowMajor); + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + coopmat adjusted = + accum_int32[i][j] - izp_bcast[i] * wsum_bcast; + coopmat adjusted_fp = + coopmat(adjusted); + coopmat scales_outer = + ifs_bcast[i] * wsc_bcast; + result[i][j] += adjusted_fp * scales_outer; + accum_int32[i][j] = coopmat(0); + } + } + } + } // groups + + // --- Bias (optional) --- +#ifdef HAS_BIAS + if (apply_bias > 0) { + for (uint t = gl_LocalInvocationID.x; t < WG_TILE_N; t += WG_SIZE) { + bias_sh[t] = float(t_bias[tile_n_start + t]); + } + memoryBarrierShared(); + barrier(); + } +#endif + + // --- Store result tile --- + // N for the store address math MUST come from the spec constant, not the + // sizes UBO (see out_N_arg above). +#ifdef IO_TEXTURE + // Epilogue iteration i drains accumulator row-block i from EVERY subgroup + // into Csh_out at once, so the SG_GRID_Y bands it holds are disjoint global + // row ranges; the whole workgroup then imageStores them. lr / MMA_M is the + // writing subgroup's warpInTile.y, so the global row reproduces the buffer + // path's gi exactly. + // + // PORTABILITY NOTE: the barrier() in the loop body keeps this loop rolled + // despite [[unroll]], so result[i][j] IS dynamically indexed. Coopmat arrays + // are opaque per-lane storage and dynamic indexing is exactly the construct + // the Xclipse/AMD-PAL compiler has broken before -- check this first if the + // texture variants miscompile on M51. + const uint CSH_TEXELS_PER_ROW = WG_TILE_N / 4u; + const uint CSH_TEXELS = CSH_ROWS * CSH_TEXELS_PER_ROW; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + // Guards Csh_out against the previous iteration's readers. Inert on i == 0 + // but must stay unconditional to remain workgroup-uniform. + barrier(); + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { +#ifdef HAS_BIAS + if (apply_bias > 0) { + const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopmat bias_tile; + coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); + result[i][j] += bias_tile; + } +#endif + coopmat out_tile = + coopmat(result[i][j]); + coopMatStore( + out_tile, Csh_out, + warpInTile.y * MMA_M * WG_TILE_N + + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j), + WG_TILE_N, + gl_CooperativeMatrixLayoutRowMajor); + } + memoryBarrierShared(); + barrier(); + + for (uint t = gl_LocalInvocationID.x; t < CSH_TEXELS; t += WG_SIZE) { + const uint lr = t / CSH_TEXELS_PER_ROW; + const uint lc4 = t % CSH_TEXELS_PER_ROW; + const uint m = + tile_m_start + (lr / MMA_M) * SG_TILE_M + i * MMA_M + (lr % MMA_M); + const uint base = lr * WG_TILE_N + lc4 * 4u; + imageStore( + t_output, + ivec3(tile_n_start / 4u + lc4, m, 0), + vec4( + float(Csh_out[base]), + float(Csh_out[base + 1u]), + float(Csh_out[base + 2u]), + float(Csh_out[base + 3u]))); + } + } +#else + const uint N_out = uint(out_N_arg); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint gi = tile_m_start + MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + const uint gj = tile_n_start + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + +#ifdef HAS_BIAS + if (apply_bias > 0) { + const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopmat bias_tile; + coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); + result[i][j] += bias_tile; + } +#endif + + coopmat out_tile = + coopmat(result[i][j]); + coopMatStore( + out_tile, t_output, + gi * N_out + gj, N_out, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#endif // IO_TEXTURE +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4.yaml b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4.yaml new file mode 100644 index 00000000000..54eba1e9c4e --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4.yaml @@ -0,0 +1,2164 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# specs/041-dbuf4-tile-sweep: TILE/SUBGROUP SWEEP variants of the int8 +# dq8ca_q4gsw coopmat shader's dbuf4 loop structure (the ORIGINAL loop before +# specs/025 User Story 1 picked dbuf2) -- +# linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4.glsl is a fork of +# linear_dq8ca_q4gsw_coopmat_tsweep.glsl (dbuf2) with only the loop structure +# swapped. Only tile geometry (WG_TILE_*, SG_GRID_*, SUBGROUP_SIZE) varies +# per variant. Selected via +# ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4_txkgs. +# Seed variant below matches the current production 8da4w tile (dbuf2, +# specs/026/027) as a legal, known-fast starting point for sweep.py's +# Optuna search; specs/041's sweep appends further candidates here as it +# runs. + +linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4: + parameter_names_with_default_values: + PRECISION: highp + HAS_BIAS: false + IO_STORAGE: buffer + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + MMA_M: 16 + MMA_N: 16 + MMA_K: 16 + WG_TILE_M: 64 + WG_TILE_N: 32 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + shader_variants: + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x32k32g12s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 32 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x32k32g12s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 32 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x256k16g82s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 256 + WG_TILE_K: 16 + SG_GRID_X: 8 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x256k16g82s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 256 + WG_TILE_K: 16 + SG_GRID_X: 8 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t256x128k64g44s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 256 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 4 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t256x128k64g44s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 256 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 4 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x64k128g44s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 64 + WG_TILE_K: 128 + SG_GRID_X: 4 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x64k128g44s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 64 + WG_TILE_K: 128 + SG_GRID_X: 4 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t256x32k64g21s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 256 + WG_TILE_N: 32 + WG_TILE_K: 64 + SG_GRID_X: 2 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t256x32k64g21s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 256 + WG_TILE_N: 32 + WG_TILE_K: 64 + SG_GRID_X: 2 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k64g42s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 64 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k64g42s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 64 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k16g41s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 4 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k16g41s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 4 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k128g84s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 128 + SG_GRID_X: 8 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k128g84s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 128 + SG_GRID_X: 8 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k16g24s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k16g24s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x256k32g24s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 256 + WG_TILE_K: 32 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x256k32g24s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 256 + WG_TILE_K: 32 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x32k64g11s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 32 + WG_TILE_K: 64 + SG_GRID_X: 1 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x32k64g11s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 32 + WG_TILE_K: 64 + SG_GRID_X: 1 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t16x128k64g21s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 16 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 2 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t16x128k64g21s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 16 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 2 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k64g81s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 8 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k64g81s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 8 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t32x256k64g22s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 32 + WG_TILE_N: 256 + WG_TILE_K: 64 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t32x256k64g22s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 32 + WG_TILE_N: 256 + WG_TILE_K: 64 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t256x128k32g12s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 256 + WG_TILE_N: 128 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t256x128k32g12s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 256 + WG_TILE_N: 128 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t256x64k16g14s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 256 + WG_TILE_N: 64 + WG_TILE_K: 16 + SG_GRID_X: 1 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t256x64k16g14s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 256 + WG_TILE_N: 64 + WG_TILE_K: 16 + SG_GRID_X: 1 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x16k32g12s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 16 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x16k32g12s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 16 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k16g24s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k16g24s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t32x32k32g12s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 32 + WG_TILE_N: 32 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t32x32k32g12s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 32 + WG_TILE_N: 32 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k64g22s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 64 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k64g22s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 64 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k16g24s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k16g24s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x32k16g12s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 32 + WG_TILE_K: 16 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x32k16g12s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 32 + WG_TILE_K: 16 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t32x128k64g11s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 32 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 1 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t32x128k64g11s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 32 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 1 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k32g41s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k32g41s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t256x128k16g81s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 256 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 8 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t256x128k16g81s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 256 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 8 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k16g28s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 2 + SG_GRID_Y: 8 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k16g28s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 2 + SG_GRID_Y: 8 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k32g44s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k32g44s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x256k16g81s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 256 + WG_TILE_K: 16 + SG_GRID_X: 8 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x256k16g81s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 256 + WG_TILE_K: 16 + SG_GRID_X: 8 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k16g12s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k16g12s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k32g11s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k32g11s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t256x128k64g82s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 256 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 8 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t256x128k64g82s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 256 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 8 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x32k128g11s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 32 + WG_TILE_K: 128 + SG_GRID_X: 1 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x32k128g11s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 32 + WG_TILE_K: 128 + SG_GRID_X: 1 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x32k64g14s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 32 + WG_TILE_K: 64 + SG_GRID_X: 1 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x32k64g14s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 32 + WG_TILE_K: 64 + SG_GRID_X: 1 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k64g24s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k64g24s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x256k16g24s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 256 + WG_TILE_K: 16 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x256k16g24s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 256 + WG_TILE_K: 16 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k64g28s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 2 + SG_GRID_Y: 8 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k64g28s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 2 + SG_GRID_Y: 8 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k64g18s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 64 + SG_GRID_X: 1 + SG_GRID_Y: 8 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k64g18s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 64 + SG_GRID_X: 1 + SG_GRID_Y: 8 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k64g42s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k64g42s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x64k16g22s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 64 + WG_TILE_K: 16 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x64k16g22s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 64 + WG_TILE_K: 16 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t32x256k16g12s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 32 + WG_TILE_N: 256 + WG_TILE_K: 16 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t32x256k16g12s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 32 + WG_TILE_N: 256 + WG_TILE_K: 16 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x256k16g28s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 256 + WG_TILE_K: 16 + SG_GRID_X: 2 + SG_GRID_Y: 8 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x256k16g28s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 256 + WG_TILE_K: 16 + SG_GRID_X: 2 + SG_GRID_Y: 8 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t16x16k32g11s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 16 + WG_TILE_N: 16 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t16x16k32g11s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 16 + WG_TILE_N: 16 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k32g14s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k32g14s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k32g12s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k32g12s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x256k32g12s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 256 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x256k32g12s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 256 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t256x128k16g22s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 256 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t256x128k16g22s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 256 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t256x32k64g12s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 256 + WG_TILE_N: 32 + WG_TILE_K: 64 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t256x32k64g12s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 256 + WG_TILE_N: 32 + WG_TILE_K: 64 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k128g42s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 128 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k128g42s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 128 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x256k64g18s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 256 + WG_TILE_K: 64 + SG_GRID_X: 1 + SG_GRID_Y: 8 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x256k64g18s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 256 + WG_TILE_K: 64 + SG_GRID_X: 1 + SG_GRID_Y: 8 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k64g14s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 1 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k64g14s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 1 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x32k64g24s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 32 + WG_TILE_K: 64 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x32k64g24s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 32 + WG_TILE_K: 64 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k64g81s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 8 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k64g81s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 8 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k128g22s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 128 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k128g22s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 128 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x64k64g12s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 64 + WG_TILE_K: 64 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x64k64g12s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 64 + WG_TILE_K: 64 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x256k64g82s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 256 + WG_TILE_K: 64 + SG_GRID_X: 8 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x256k64g82s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 256 + WG_TILE_K: 64 + SG_GRID_X: 8 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t32x64k64g42s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 32 + WG_TILE_N: 64 + WG_TILE_K: 64 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t32x64k64g42s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 32 + WG_TILE_N: 64 + WG_TILE_K: 64 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x32k16g12s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 32 + WG_TILE_K: 16 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x32k16g12s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 32 + WG_TILE_K: 16 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k16g42s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k16g42s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t256x128k32g48s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 256 + WG_TILE_N: 128 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 8 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t256x128k32g48s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 256 + WG_TILE_N: 128 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 8 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k32g44s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k32g44s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k16g18s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 1 + SG_GRID_Y: 8 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k16g18s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 1 + SG_GRID_Y: 8 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t32x128k16g21s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 32 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 2 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t32x128k16g21s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 32 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 2 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t16x128k64g81s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 16 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 8 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t16x128k64g81s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 16 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 8 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k16g18s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 16 + SG_GRID_X: 1 + SG_GRID_Y: 8 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k16g18s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 16 + SG_GRID_X: 1 + SG_GRID_Y: 8 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k64g22s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 64 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k64g22s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 64 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k32g24s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 32 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k32g24s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 32 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x32k64g22s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 32 + WG_TILE_K: 64 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x32k64g22s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 32 + WG_TILE_K: 64 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x32k64g28s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 32 + WG_TILE_K: 64 + SG_GRID_X: 2 + SG_GRID_Y: 8 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x32k64g28s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 32 + WG_TILE_K: 64 + SG_GRID_X: 2 + SG_GRID_Y: 8 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k32g41s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k32g41s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k128g28s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 128 + SG_GRID_X: 2 + SG_GRID_Y: 8 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k128g28s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 128 + SG_GRID_X: 2 + SG_GRID_Y: 8 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t256x32k16g11s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 256 + WG_TILE_N: 32 + WG_TILE_K: 16 + SG_GRID_X: 1 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t256x32k16g11s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 256 + WG_TILE_N: 32 + WG_TILE_K: 16 + SG_GRID_X: 1 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k64g44s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 4 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k64g44s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 4 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k16g14s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 16 + SG_GRID_X: 1 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k16g14s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 16 + SG_GRID_X: 1 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x16k64g12s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 16 + WG_TILE_K: 64 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x16k64g12s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 16 + WG_TILE_K: 64 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k16g21s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 16 + SG_GRID_X: 2 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k16g21s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 16 + SG_GRID_X: 2 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k32g18s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 8 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k32g18s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 8 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k64g24s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k64g24s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k32g42s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k32g42s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k128g18s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 128 + SG_GRID_X: 1 + SG_GRID_Y: 8 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k128g18s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 128 + SG_GRID_X: 1 + SG_GRID_Y: 8 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x16k64g18s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 16 + WG_TILE_K: 64 + SG_GRID_X: 1 + SG_GRID_Y: 8 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x16k64g18s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 16 + WG_TILE_K: 64 + SG_GRID_X: 1 + SG_GRID_Y: 8 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k64g21s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 2 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k64g21s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 2 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k32g24s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 32 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k32g24s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 32 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x256k32g84s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 256 + WG_TILE_K: 32 + SG_GRID_X: 8 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x256k32g84s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 256 + WG_TILE_K: 32 + SG_GRID_X: 8 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k16g44s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 4 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k16g44s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 4 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k32g12s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k32g12s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k64g11s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 64 + SG_GRID_X: 1 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k64g11s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 64 + SG_GRID_X: 1 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x16k64g12s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 16 + WG_TILE_K: 64 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x16k64g12s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 16 + WG_TILE_K: 64 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t256x128k64g24s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 256 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t256x128k64g24s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 256 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k32g24s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k32g24s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x32k128g12s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 32 + WG_TILE_K: 128 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x32k128g12s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 32 + WG_TILE_K: 128 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k128g44s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 128 + SG_GRID_X: 4 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k128g44s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 128 + SG_GRID_X: 4 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k32g24s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 32 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k32g24s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 32 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k64g84s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 8 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k64g84s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 8 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t32x16k32g12s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 32 + WG_TILE_N: 16 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t32x16k32g12s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 32 + WG_TILE_N: 16 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x64k32g42s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x64k32g42s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t256x128k64g81s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 256 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 8 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t256x128k64g81s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 256 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 8 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k64g44s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 64 + SG_GRID_X: 4 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k64g44s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 64 + SG_GRID_X: 4 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x64k32g24s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x64k32g24s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x256k16g14s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 256 + WG_TILE_K: 16 + SG_GRID_X: 1 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x256k16g14s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 256 + WG_TILE_K: 16 + SG_GRID_X: 1 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x256k32g44s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 256 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x256k32g44s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 256 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x32k32g24s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 32 + WG_TILE_K: 32 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x32k32g24s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 32 + WG_TILE_K: 32 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k64g24s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 64 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k64g24s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 64 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k32g14s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k32g14s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x32k16g12s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 32 + WG_TILE_K: 16 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x32k16g12s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 32 + WG_TILE_K: 16 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k32g14s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k32g14s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k32g24s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 32 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k32g24s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 32 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k16g24s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 16 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k16g24s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 16 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x256k16g12s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 256 + WG_TILE_K: 16 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x256k16g12s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 256 + WG_TILE_K: 16 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x256k16g24s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 256 + WG_TILE_K: 16 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x256k16g24s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 256 + WG_TILE_K: 16 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k64g28s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 2 + SG_GRID_Y: 8 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k64g28s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 2 + SG_GRID_Y: 8 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k16g14s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 1 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k16g14s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 1 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x256k32g18s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 256 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 8 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x256k32g18s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 256 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 8 + SUBGROUP_SIZE: 32 diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_qw_coopmat.glsl b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_qw_coopmat.glsl index 755261452f4..5edb05001c7 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_qw_coopmat.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_qw_coopmat.glsl @@ -7,6 +7,15 @@ */ /* + * specs/025/026/027: loop structure updated from "dbuf4" (store-first, + * single-buffered-until-prefetch) to "dbuf2" (store-first, prefetch-first + * prologue) per specs/025 User Story 1's re-confirmed loop-structure winner + * for this shader; tile geometry updated from the prior 128x64/K32/2x2/s64 + * to 64x32/K32/1x2/s64 per specs/027's e2e-ranked sweep winner + * (tsweep_t64x32k32g12s64) -- confirmed +9.32% real end-to-end prefill + * throughput on M5 EVT1 (Llama 3.1 8B, 2048-token prefill), not just + * isolated-kernel GFLOP/s. See specs/027-e2e-tile-sweep/results/sweep-report.md. + * * KHR Cooperative Matrix variant of the dynamically-quantized-activation * linear tiled shader (WEIGHT_NBITS=4): * 4 -> linear_dq8ca_q4gsw_coopmat INT4 group-symmetric weight @@ -25,15 +34,20 @@ * "8 * input_sum" term of the tiled correction (which compensates for * unsigned int4 nibbles in dotPacked4x8) cancels out and is not needed. * - * Loop structure follows the NVIDIA double-buffered GEMM reference - * (shmem_double_buf4.comp "store-first" variant; see coopmat_mm_ref.glsl in - * test/custom_ops): prologue register prefetch, then per chunk - * barrier -> prefetch next chunk -> int8 MMA on the current LDS slice -> - * store temp into the other slice. One barrier per chunk; the prefetch is - * pure loads, in flight during the math; quant unpack happens at the store - * stage. The loop stays NESTED (groups x chunks, group epilog unconditional - * at the group tail) — flattening it with a conditional coopmat epilog - * crashes the Xclipse PAL compiler at large spec-resolved trip counts. + * Loop structure ("dbuf2", specs/023-8da4w-int8-dbuf-sweep naming): prologue + * prefetches chunk 0 into temp registers only (no shared-memory write, no + * barrier); each loop iteration then does store(temp -> cur slice) + * -> barrier() [UNCONDITIONAL, every iteration] -> MMA(cur) -> prefetch(next + * chunk -> temp) [skipped on the last chunk]. Iteration `chunk` stores the + * data FOR ITSELF (already prefetched by the previous iteration, or by the + * prologue for chunk 0), immediately before using it. The same inversion + * applies to the group wsum/wsc ping-pong: this variant stores the CURRENT + * group's values (prefetched by the previous group's last chunk) at the head + * of the group's first chunk. Group 0's wsum/wsc are unaffected -- set up + * directly in the prologue. The nested groups x chunks loop and + * unconditional group epilog are kept exactly as before -- flattening them + * with a conditional coopmat epilog crashes the Xclipse PAL compiler at + * large spec-resolved trip counts (specs/023 finding). * * Per-(group, N) weight sums/scales live in a SECOND ping-pong pair indexed * by group parity: the next group's values are prefetched into registers @@ -49,11 +63,13 @@ * per lane with a bank-conflict-free col stride. Each uint holds 4 packed * int8. * - * Tile hierarchy (yaml): MMA 16x16x16 int8, WG_TILE 128x64, WG_TILE_K = 32, - * 4 subgroups x 64 threads. The double-buffered reference's subgroup-32 - * layout is NOT used: the Xclipse PAL compiler crashes in - * vkCreateComputePipelines when int8 WMMA is compiled at forced subgroup - * size 32 (fp16 WMMA at 32 is fine; see linear_qw_coopmat). + * Tile hierarchy (yaml): MMA 16x16x16 int8, WG_TILE 64x32, WG_TILE_K = 32, + * 2 subgroups x 64 threads (1x2 grid) -- specs/027's e2e-ranked winner. + * SUBGROUP_SIZE stays 64: specs/026 found subgroup=32 is legal (no compiler + * crash) but sharply tile-shape-dependently INCORRECT, and this tile shape + * was not one of the two shapes specs/026 found fully-correct at subgroup=32 + * -- see specs/026-8da4w-subgroup32-sweep/results/ for the full picture + * before considering subgroup=32 at this or any other tile shape. * * Hard preconditions: * M % WG_TILE_M == 0, N % WG_TILE_N == 0, K % WG_TILE_K == 0, @@ -114,6 +130,12 @@ ${layout_declare_spec_const(C, "int", "K4_per_group", "0")} // crashes (null deref in vkCreateComputePipelines) when a loop containing // coopMatMulAdd has a UBO-derived trip count. INT4: number of quant groups; // INT8: number of K-chunks. +// +// Unlike linear_qw_coopmat, this spec-const workaround is INTENTIONALLY kept +// here: on 2026-06-30 the UBO-direct method (sizes UBO feeding num_chunks/N +// directly) was A/B'd on this shader and produced wrong results for the +// coopmat (buffer) path at M>=128, while this spec-const version validated +// clean — see add_linear_dqa_qw_node in QuantizedLinear.cpp. ${layout_declare_spec_const(C, "int", "num_groups_arg", "0")} // Output width N for coopMatStore: the Xclipse compiler MISCOMPILES // coopMatStore whose offset/stride derive from a UBO value (only the first @@ -297,8 +319,10 @@ void main() { gl_CooperativeMatrixLayoutColumnMajor); } - // Prefetch chunk 0 into temp registers, then store to slice 0 (no barrier; - // the first loop iteration's barrier publishes it). + // dbuf2: prefetch chunk 0 into temp registers only -- no shared-memory + // write, no barrier here. The main loop's first iteration stores temp + // into slice 0 and barriers as normal (uniform code path for every chunk, + // including chunk 0). if (a_active) { const uint m4_global = (tile_m_start >> 2u) + a_m_block; temp_A = t_packed_int8_input[m4_global * nblocks_x_A + a_k_block]; @@ -325,116 +349,88 @@ void main() { #endif } #endif - { - // store chunk 0 -> slice 0 - if (a_active) { - const uint slab_idx = a_k_block / (MMA_K >> 2u); - const uint k_uint_in_slab = a_k_block % (MMA_K >> 2u); - const uint base_row = a_m_block * 4u; - [[unroll]] for (uint m4i = 0; m4i < 4u; ++m4i) { - Ash_int8[slab_idx * A_SLAB_U32 + (base_row + m4i) * A_STRIDE_U32 + k_uint_in_slab] = - uint(temp_A[m4i]); - } - } -#ifdef WEIGHT_INT4 - [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { - const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; - const uint block_in_chunk = slot >> 3u; - const uint col_in_block = slot & 7u; - const uint k4_in_chunk = block_in_chunk / N8_PER_TILE; - const uint n8_in_tile = block_in_chunk % N8_PER_TILE; - const uint r = col_in_block & 3u; - const uint parity = col_in_block >> 2u; - const int w = temp_B[si][r]; - const int base = int(4u * parity); - const int v0 = (((w >> (base + 0)) & 0xF) - 8) & 0xFF; - const int v1 = (((w >> (base + 8)) & 0xF) - 8) & 0xFF; - const int v2 = (((w >> (base + 16)) & 0xF) - 8) & 0xFF; - const int v3 = (((w >> (base + 24)) & 0xF) - 8) & 0xFF; - const uint n_col = n8_in_tile * 8u + r + parity * 4u; - const uint slab_idx = k4_in_chunk / (MMA_K >> 2u); - const uint k4_in_slab = k4_in_chunk % (MMA_K >> 2u); - Bsh_int8[slab_idx * B_SLAB_U32 + n_col * B_STRIDE_U32 + k4_in_slab] = - uint(v0 | (v1 << 8) | (v2 << 16) | (v3 << 24)); - } -#else - if (b_active) { - const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); - const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); - const uint n_col_base = b_n_uint_col * 4u; - [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { - Bsh_int8[slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab] = - uint(temp_B[n_in_blk]); - } - } -#endif - } // ========================================================= // MAIN LOOP — nested groups x chunks (the flattened single loop with a // conditional coopmat epilog crashes the Xclipse PAL compiler at large - // spec-resolved trip counts). One barrier per chunk. Chunk iteration - // (global index `chunk`): - // 1. barrier — A/B slice (chunk%2) fully written; on the first chunk - // of group g, wsum/wsc slice (g%2) is too. - // 2. prefetch — chunk+1 (A blocks, B blocks) into temp; when chunk+1 - // starts a new group, also its wsum/wsc element. Skipped - // entirely on the final chunk. + // spec-resolved trip counts). One barrier per chunk, UNCONDITIONAL every + // iteration (dbuf2 "store-first" ordering). Chunk iteration (global index + // `chunk`): + // 1. store — temp (already holding this chunk's data, from the + // prologue for chunk 0 or from the previous iteration's + // step 4) -> A/B slice (chunk%2), unpacking the weight; + // on a group boundary (first chunk of a group > 0), + // also store this group's wsum/wsc -> slice (group_i%2). + // 2. barrier — A/B slice (chunk%2) (and, on a group boundary, wsum/wsc + // slice (group_i%2)) fully written; UNCONDITIONAL, not + // skipped on the last chunk. // 3. int8 MMA — on slice (chunk%2) into accum_int32. - // 4. store — temp -> A/B slice ((chunk+1)%2), unpacking the weight; - // on a group boundary, wsum/wsc -> slice ((g+1)%2). + // 4. prefetch — chunk+1 (A blocks, B blocks) into temp; when chunk+1 + // starts a new group, also that group's wsum/wsc + // element. Skipped entirely on the final chunk. // The group epilog runs unconditionally at the tail of each group. // ========================================================= uint chunk = 0; for (uint group_i = 0; group_i < num_groups; ++group_i) { for (uint inner = 0; inner < CHUNKS_PER_GROUP; ++inner, ++chunk) { const bool has_next = chunk + 1u < num_chunks; - const bool group_crossing = has_next && (inner + 1u == CHUNKS_PER_GROUP); const uint cur_a = (chunk % 2u) * ASH_SLICE_U32; const uint cur_b = (chunk % 2u) * BSH_SLICE_U32; - const uint nxt_a = ((chunk + 1u) % 2u) * ASH_SLICE_U32; - const uint nxt_b = ((chunk + 1u) % 2u) * BSH_SLICE_U32; - barrier(); - - // --- 2. prefetch chunk+1 -> temp --- - if (has_next) { - const uint chunkK_nxt = (chunk + 1u) * WG_TILE_K; - if (a_active) { - const uint m4_global = (tile_m_start >> 2u) + a_m_block; - const uint k4_global = (chunkK_nxt >> 2u) + a_k_block; - temp_A = t_packed_int8_input[m4_global * nblocks_x_A + k4_global]; + // --- 1. store temp (this chunk) -> cur slice --- + if (a_active) { + const uint slab_idx = a_k_block / (MMA_K >> 2u); + const uint k_uint_in_slab = a_k_block % (MMA_K >> 2u); + const uint base_row = a_m_block * 4u; + [[unroll]] for (uint m4i = 0; m4i < 4u; ++m4i) { + Ash_int8[cur_a + slab_idx * A_SLAB_U32 + (base_row + m4i) * A_STRIDE_U32 + k_uint_in_slab] = + uint(temp_A[m4i]); } + } #ifdef WEIGHT_INT4 - [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { - const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; - const uint block_in_chunk = slot >> 3u; - const uint k4_blk = (chunkK_nxt >> 2u) + block_in_chunk / N8_PER_TILE; - const uint n8_blk = (tile_n_start >> 3u) + (block_in_chunk % N8_PER_TILE); -#ifdef WEIGHT_BUFFER - temp_B[si] = t_packed_weight[(n8_blk * nblocks_x_A) + k4_blk]; -#else - temp_B[si] = texelFetch(t_packed_weight, ivec2(k4_blk, n8_blk), 0); -#endif - } - if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { - const uint n_idx = tile_n_start + gl_LocalInvocationID.x; - f16vec4 sv = t_weight_scales[(group_i + 1u) * N4 + (n_idx >> 2u)]; - temp_wsc = float(sv[n_idx & 3u]); - temp_wsum = t_weight_sums[(group_i + 1u) * N + n_idx]; - } -#else - if (b_active) { - const uint block_y_w = (chunkK_nxt >> 2u) + b_k4_in_chunk; - const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; -#ifdef WEIGHT_BUFFER - temp_B = t_packed_weight[(block_y_w * N4) + block_x_w]; + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { + const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; + const uint block_in_chunk = slot >> 3u; + const uint col_in_block = slot & 7u; + const uint k4_in_chunk = block_in_chunk / N8_PER_TILE; + const uint n8_in_tile = block_in_chunk % N8_PER_TILE; + const uint r = col_in_block & 3u; + const uint parity = col_in_block >> 2u; + const int w = temp_B[si][r]; + const int base = int(4u * parity); + const int v0 = (((w >> (base + 0)) & 0xF) - 8) & 0xFF; + const int v1 = (((w >> (base + 8)) & 0xF) - 8) & 0xFF; + const int v2 = (((w >> (base + 16)) & 0xF) - 8) & 0xFF; + const int v3 = (((w >> (base + 24)) & 0xF) - 8) & 0xFF; + const uint n_col = n8_in_tile * 8u + r + parity * 4u; + const uint slab_idx = k4_in_chunk / (MMA_K >> 2u); + const uint k4_in_slab = k4_in_chunk % (MMA_K >> 2u); + Bsh_int8[cur_b + slab_idx * B_SLAB_U32 + n_col * B_STRIDE_U32 + k4_in_slab] = + uint(v0 | (v1 << 8) | (v2 << 16) | (v3 << 24)); + } + // Group boundary (this is the first chunk of a group other than + // group 0): store this group's wsum/wsc, prefetched by the previous + // group's last chunk (step 4 below). + if (inner == 0u && group_i > 0u && gl_LocalInvocationID.x < WG_TILE_N) { + const uint wbase_cur = (group_i % 2u) * WG_TILE_N; + wsum_sh[wbase_cur + gl_LocalInvocationID.x] = temp_wsum; + wsc_sh[wbase_cur + gl_LocalInvocationID.x] = temp_wsc; + } #else - temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, block_y_w), 0); -#endif + if (b_active) { + const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); + const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); + const uint n_col_base = b_n_uint_col * 4u; + [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { + Bsh_int8[cur_b + slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab] = + uint(temp_B[n_in_blk]); } -#endif } +#endif + + // --- 2. barrier — cur slice(s) fully written --- + memoryBarrierShared(); + barrier(); // --- 3. int8 MMA on the cur slice --- [[unroll]] for (uint k = 0; k < NUM_K_SLABS; ++k) { @@ -465,52 +461,42 @@ void main() { } } - // --- 4. store temp (chunk+1) -> nxt slice --- + // --- 4. prefetch chunk+1 -> temp --- if (has_next) { + const bool group_crossing = (inner + 1u == CHUNKS_PER_GROUP); + const uint chunkK_nxt = (chunk + 1u) * WG_TILE_K; if (a_active) { - const uint slab_idx = a_k_block / (MMA_K >> 2u); - const uint k_uint_in_slab = a_k_block % (MMA_K >> 2u); - const uint base_row = a_m_block * 4u; - [[unroll]] for (uint m4i = 0; m4i < 4u; ++m4i) { - Ash_int8[nxt_a + slab_idx * A_SLAB_U32 + (base_row + m4i) * A_STRIDE_U32 + k_uint_in_slab] = - uint(temp_A[m4i]); - } + const uint m4_global = (tile_m_start >> 2u) + a_m_block; + const uint k4_global = (chunkK_nxt >> 2u) + a_k_block; + temp_A = t_packed_int8_input[m4_global * nblocks_x_A + k4_global]; } #ifdef WEIGHT_INT4 [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; const uint block_in_chunk = slot >> 3u; - const uint col_in_block = slot & 7u; - const uint k4_in_chunk = block_in_chunk / N8_PER_TILE; - const uint n8_in_tile = block_in_chunk % N8_PER_TILE; - const uint r = col_in_block & 3u; - const uint parity = col_in_block >> 2u; - const int w = temp_B[si][r]; - const int base = int(4u * parity); - const int v0 = (((w >> (base + 0)) & 0xF) - 8) & 0xFF; - const int v1 = (((w >> (base + 8)) & 0xF) - 8) & 0xFF; - const int v2 = (((w >> (base + 16)) & 0xF) - 8) & 0xFF; - const int v3 = (((w >> (base + 24)) & 0xF) - 8) & 0xFF; - const uint n_col = n8_in_tile * 8u + r + parity * 4u; - const uint slab_idx = k4_in_chunk / (MMA_K >> 2u); - const uint k4_in_slab = k4_in_chunk % (MMA_K >> 2u); - Bsh_int8[nxt_b + slab_idx * B_SLAB_U32 + n_col * B_STRIDE_U32 + k4_in_slab] = - uint(v0 | (v1 << 8) | (v2 << 16) | (v3 << 24)); + const uint k4_blk = (chunkK_nxt >> 2u) + block_in_chunk / N8_PER_TILE; + const uint n8_blk = (tile_n_start >> 3u) + (block_in_chunk % N8_PER_TILE); +#ifdef WEIGHT_BUFFER + temp_B[si] = t_packed_weight[(n8_blk * nblocks_x_A) + k4_blk]; +#else + temp_B[si] = texelFetch(t_packed_weight, ivec2(k4_blk, n8_blk), 0); +#endif } if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { - const uint wbase_nxt = ((group_i + 1u) % 2u) * WG_TILE_N; - wsum_sh[wbase_nxt + gl_LocalInvocationID.x] = temp_wsum; - wsc_sh[wbase_nxt + gl_LocalInvocationID.x] = temp_wsc; + const uint n_idx = tile_n_start + gl_LocalInvocationID.x; + f16vec4 sv = t_weight_scales[(group_i + 1u) * N4 + (n_idx >> 2u)]; + temp_wsc = float(sv[n_idx & 3u]); + temp_wsum = t_weight_sums[(group_i + 1u) * N + n_idx]; } #else if (b_active) { - const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); - const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); - const uint n_col_base = b_n_uint_col * 4u; - [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { - Bsh_int8[nxt_b + slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab] = - uint(temp_B[n_in_blk]); - } + const uint block_y_w = (chunkK_nxt >> 2u) + b_k4_in_chunk; + const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; +#ifdef WEIGHT_BUFFER + temp_B = t_packed_weight[(block_y_w * N4) + block_x_w]; +#else + temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, block_y_w), 0); +#endif } #endif } diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_qw_coopmat.yaml b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_qw_coopmat.yaml index 959cb51966d..7bc39d224a8 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_qw_coopmat.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_qw_coopmat.yaml @@ -10,11 +10,13 @@ # WEIGHT_NBITS=4 -> linear_dq8ca_q4gsw_coopmat (INT4 group-symmetric) # Requires the VK_COMPONENT_TYPE_SINT8_KHR cooperative matrix property to be # enumerated on the device. -# Loop structure follows the double-buffered reference (coopmat_mm_ref) at -# a 128x64 tile with K-step 32, 4 subgroups x 64 threads. The reference's -# subgroup-32 layout is NOT used — the Xclipse PAL compiler crashes in -# vkCreateComputePipelines when int8 WMMA is compiled at forced subgroup -# size 32 (fp16 WMMA at 32 is fine; see linear_qw_coopmat). +# specs/027-e2e-tile-sweep: tile geometry updated from 128x64/K32/2x2 to +# 64x32/K32/1x2 (specs/027's e2e-ranked sweep winner, +9.32% confirmed real +# end-to-end prefill throughput vs the prior geometry, not just isolated +# microbenchmark GFLOP/s). SUBGROUP_SIZE stays 64: specs/026 found +# subgroup=32 is legal (no compiler crash) but sharply tile-shape-dependently +# incorrect, and this tile shape was not verified correct at subgroup=32 -- +# see specs/026-8da4w-subgroup32-sweep/results/ before considering it. linear_dq8ca_qw_coopmat: parameter_names_with_default_values: @@ -25,10 +27,10 @@ linear_dq8ca_qw_coopmat: MMA_M: 16 MMA_N: 16 MMA_K: 16 - WG_TILE_M: 128 - WG_TILE_N: 64 + WG_TILE_M: 64 + WG_TILE_N: 32 WG_TILE_K: 32 - SG_GRID_X: 2 + SG_GRID_X: 1 SG_GRID_Y: 2 SUBGROUP_SIZE: 64 shader_variants: diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_coopmat_tsweep_dbuf4.glsl b/backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_coopmat_tsweep_dbuf4.glsl new file mode 100644 index 00000000000..033afc969e8 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_coopmat_tsweep_dbuf4.glsl @@ -0,0 +1,473 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +/* + * TILE/SUBGROUP-SWEEP variant of linear_q4gsw_coopmat's dbuf4 ("store-first", + * prefetch-peeled) loop structure (specs/041-dbuf4-tile-sweep). Forked from + * linear_q4gsw_coopmat_tsweep.glsl (which carries dbuf1's loop, the + * production winner) -- everything except the PROLOGUE/MAIN LOOP block is + * identical: same bindings, spec-constants, tile-geometry templating, + * dequant_block, fp16 accumulator, bias/store epilogue. Only the loop + * structure is swapped to dbuf4, recovered from the pre-retune reference + * (git commit 8d0f23ee78's sibling; see specs/041/reference/) that predates + * the dbuf1-winner fp16-accumulate/flattened-loop retune. dbuf4: + * prologue: prefetch chunk 0 -> temp, store to slice 0 (no barrier) + * per iter: barrier -> prefetch(next) -> MMA(cur) -> store(next) + * epilogue: barrier -> MMA(last) [loop peeled, 1 barrier/iter] + * vs. dbuf1's single flattened loop with a trailing conditional barrier + * (skipped on the last chunk) -- dbuf4 issues one barrier per chunk + * (num_chunks total) instead of dbuf1's (num_chunks - 1). + * + * Selected at dispatch via + * ET_VK_Q4GSW_COOPMAT_VARIANT=tsweep_dbuf4_txkgs + * (QuantizedLinear.cpp), additive to the existing tsweep_t... (dbuf1) + * namespace. + * + * Hard preconditions (no shape/alignment checks inside the shader): + * M % WG_TILE_M == 0 + * N % WG_TILE_N == 0 + * K % WG_TILE_K == 0 + * group_size % WG_TILE_K == 0 + * Misaligned shapes silently miscompute / overrun -- gated at dispatch time + * by can_use_q4gsw_coopmat() using the ACTIVE variant's own tile dims. + */ + +#version 450 core + +#extension GL_KHR_cooperative_matrix : require +#extension GL_KHR_memory_scope_semantics : require +#extension GL_KHR_shader_subgroup_basic : enable +#extension GL_EXT_shader_explicit_arithmetic_types : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_EXT_control_flow_attributes : enable + +#define PRECISION ${PRECISION} + +$if HAS_BIAS: + #define HAS_BIAS + +$if WEIGHT_STORAGE == "buffer": + #define WEIGHT_BUFFER + +$if IO_STORAGE == "texture3d": + #define IO_TEXTURE + +layout(std430) buffer; + +#include "common.glslh" + +${layout_declare_tensor(B, "w", "t_output", "half", IO_STORAGE, is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_input", "half", IO_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_packed_weight", "int", WEIGHT_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_weight_scales", "half", "buffer", is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_bias", "half", "buffer", is_scalar_array=True)} + +${layout_declare_ubo(B, "ivec4", "output_sizes")} +${layout_declare_ubo(B, "ivec4", "input_sizes")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +${layout_declare_spec_const(C, "int", "apply_bias", "0")} +${layout_declare_spec_const(C, "int", "K4_per_group", "0")} +${layout_declare_spec_const(C, "int", "num_groups_arg", "0")} +${layout_declare_spec_const(C, "int", "out_N_arg", "0")} + +// --- Tile geometry (from yaml; per-variant tile-sweep candidate) --- +const uint MMA_M = ${MMA_M}; +const uint MMA_N = ${MMA_N}; +const uint MMA_K = ${MMA_K}; + +const uint WG_TILE_M = ${WG_TILE_M}; +const uint WG_TILE_N = ${WG_TILE_N}; +const uint WG_TILE_K = ${WG_TILE_K}; + +const uint SG_GRID_X = ${SG_GRID_X}; +const uint SG_GRID_Y = ${SG_GRID_Y}; +const uint SUBGROUP_SIZE = ${SUBGROUP_SIZE}; +const uint NUM_SUBGROUPS = SG_GRID_X * SG_GRID_Y; +const uint WG_SIZE = NUM_SUBGROUPS * SUBGROUP_SIZE; + +const uint SG_TILE_M = WG_TILE_M / SG_GRID_Y; +const uint SG_TILE_N = WG_TILE_N / SG_GRID_X; +const uint MMAS_PER_SG_M = SG_TILE_M / MMA_M; +const uint MMAS_PER_SG_N = SG_TILE_N / MMA_N; + +// fp16: 8 elements per uvec4 (128-bit) +const uint FP16_PER_VEC4 = 8; +const uint A_STRIDE_VEC4 = (WG_TILE_K + FP16_PER_VEC4) / FP16_PER_VEC4; +const uint B_STRIDE_VEC4 = (WG_TILE_N + FP16_PER_VEC4) / FP16_PER_VEC4; + +// One ping-pong slice of each shared-memory buffer (in uvec4 units). +const uint ASH_SLICE = WG_TILE_M * A_STRIDE_VEC4; +const uint BSH_SLICE = WG_TILE_K * B_STRIDE_VEC4; + +// Double-buffered shared memory. +shared uvec4 Ash[2 * ASH_SLICE]; +shared uvec4 Bsh[2 * BSH_SLICE]; +#ifdef HAS_BIAS +shared float16_t bias_sh[WG_TILE_N]; +#endif + +#ifdef IO_TEXTURE +// Result staging for the imageStore epilogue: SG_GRID_Y bands of MMA_M rows, +// each WG_TILE_N wide, row-major. A full WG_TILE_M x WG_TILE_N buffer would +// cost SG_GRID_Y/MMAS_PER_SG_M x more LDS and wreck occupancy. Not aliased +// with Ash/Bsh: GLSL has no shared unions and coopMatStore needs a +// float16_t-typed array. +const uint CSH_ROWS = SG_GRID_Y * MMA_M; +shared float16_t Csh[CSH_ROWS * WG_TILE_N]; +#endif + +// Staging thread maps: each thread covers one uvec4 (8 fp16) per pass. +const uint INVS_PER_ROW_A = WG_TILE_K / FP16_PER_VEC4; +const uint A_ROWS_PER_PASS = WG_SIZE / INVS_PER_ROW_A; +const uint A_PASSES = WG_TILE_M / A_ROWS_PER_PASS; +const uint INVS_PER_ROW_B = WG_TILE_N / FP16_PER_VEC4; +const uint B_ROWS_PER_PASS = WG_SIZE / INVS_PER_ROW_B; +const uint B_PASSES = WG_TILE_K / B_ROWS_PER_PASS; + +// FP16 accumulator coopmats (MMAS_PER_SG_M x MMAS_PER_SG_N per thread). +coopmat + result[MMAS_PER_SG_M][MMAS_PER_SG_N]; + +// Dequant one packed INT4 block column-pair into 8 scaled fp16 weights +// (one Bsh uvec4). col_lo/col_hi select the K row within the block. +uvec4 dequant_block( + const ivec4 wb, + const uint col_lo, + const uint col_hi, + const f16vec4 s0, + const f16vec4 s1) { + const f16vec4 v0 = f16vec4(((wb >> int(4u * col_lo)) & 0xF) - 8) * s0; + const f16vec4 v1 = f16vec4(((wb >> int(4u * col_hi)) & 0xF) - 8) * s1; + return uvec4( + packFloat2x16(v0.xy), packFloat2x16(v0.zw), + packFloat2x16(v1.xy), packFloat2x16(v1.zw)); +} + +// Fetch 8 consecutive fp16 activations of `row` starting at half-vec4 index +// `k_hv4`, packed into one uvec4. Both spellings address the same bytes: a +// width-packed texture3d holds elements [4x, 4x+3] of row m at texel (x, m, 0), +// so buffer index (row * K4 + x) and texel coord (x, row, 0) are one address in +// two notations. +uvec4 load_a_vec4(const uint row, const uint k_hv4, const uint K4) { +#ifdef IO_TEXTURE + // Narrow to f16vec4 AT the fetch: a half sampler is typed to return vec4 + // (fp32), and consuming the fp32 costs an image_load into 4 VGPRs plus a + // v_cvt_pk_rtz per pair. Taking it as f16vec4 immediately lets the compiler + // fold the narrowing into `image_load ... d16` (2 VGPRs, no conversion). + // Lossless -- the source is rgba16f. + const f16vec4 v0 = f16vec4(texelFetch(t_input, ivec3(k_hv4, row, 0), 0)); + const f16vec4 v1 = f16vec4(texelFetch(t_input, ivec3(k_hv4 + 1u, row, 0), 0)); + return uvec4( + packFloat2x16(v0.xy), packFloat2x16(v0.zw), + packFloat2x16(v1.xy), packFloat2x16(v1.zw)); +#else + const f16vec4 v0 = t_input[row * K4 + k_hv4]; + const f16vec4 v1 = t_input[row * K4 + k_hv4 + 1u]; + return uvec4( + packFloat2x16(v0.xy), packFloat2x16(v0.zw), + packFloat2x16(v1.xy), packFloat2x16(v1.zw)); +#endif +} + +void main() { + const uvec2 tileID = uvec2(gl_WorkGroupID.xy); + const uvec2 warpInTile = uvec2( + gl_SubgroupID % SG_GRID_X, + gl_SubgroupID / SG_GRID_X); + + const uint K = uint(input_sizes.x); + const uint K4 = (K + 3u) / 4u; + const uint N4 = (uint(output_sizes.x) + 3u) / 4u; + + const uint CHUNKS_PER_GROUP = uint(K4_per_group) * 4u / WG_TILE_K; + const uint num_chunks = uint(num_groups_arg) * CHUNKS_PER_GROUP; + + const uint tile_m_start = WG_TILE_M * tileID.y; + const uint tile_n_start = WG_TILE_N * tileID.x; + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + result[i][j] = coopmat(0.0); + } + } + + const uint a_col = gl_LocalInvocationID.x % INVS_PER_ROW_A; + const uint a_row_offset = gl_LocalInvocationID.x / INVS_PER_ROW_A; + const uint b_col = gl_LocalInvocationID.x % INVS_PER_ROW_B; + const uint b_row_offset = gl_LocalInvocationID.x / INVS_PER_ROW_B; + + const uint n8_blk = (tile_n_start + b_col * 8u) >> 3u; + const uint col_lo = 2u * (b_row_offset & 3u); + const uint col_hi = col_lo + 1u; + + const uint sc_n4 = (tile_n_start + b_col * 8u) >> 2u; + uint cached_group = 0xFFFFFFFFu; + f16vec4 sc0; + f16vec4 sc1; + + uvec4 temp_A[A_PASSES]; + ivec4 temp_B[B_PASSES]; + + // ========================================================= + // PROLOGUE (dbuf4): prefetch chunk 0 into temp registers, then store to + // slice 0. No barrier here -- the main loop's first iteration barriers + // before reading slice 0. + // ========================================================= + { + [[unroll]] for (uint p = 0; p < A_PASSES; ++p) { + const uint row = tile_m_start + p * A_ROWS_PER_PASS + a_row_offset; + const uint k_hv4 = (a_col * FP16_PER_VEC4) / 4u; + temp_A[p] = load_a_vec4(row, k_hv4, K4); + } + cached_group = 0u; + sc0 = t_weight_scales[sc_n4]; + sc1 = t_weight_scales[sc_n4 + 1u]; + [[unroll]] for (uint p = 0; p < B_PASSES; ++p) { + const uint k_row = p * B_ROWS_PER_PASS + b_row_offset; + ivec4 wblock; +#ifdef WEIGHT_BUFFER + wblock = t_packed_weight[n8_blk * K4 + (k_row >> 2u)]; +#else + wblock = texelFetch(t_packed_weight, ivec2(k_row >> 2u, n8_blk), 0); +#endif + temp_B[p] = wblock; + } + } + { + [[unroll]] for (uint p = 0; p < A_PASSES; ++p) { + Ash[(p * A_ROWS_PER_PASS + a_row_offset) * A_STRIDE_VEC4 + a_col] = temp_A[p]; + } + [[unroll]] for (uint p = 0; p < B_PASSES; ++p) { + Bsh[(p * B_ROWS_PER_PASS + b_row_offset) * B_STRIDE_VEC4 + b_col] = + dequant_block(temp_B[p], col_lo, col_hi, sc0, sc1); + } + } + + // ========================================================= + // MAIN LOOP (dbuf4) — one barrier per iteration, unconditional body, loop + // peeled (bound excludes the last chunk). Iteration `chunk` does: + // 1. barrier — slice (chunk%2) fully written + // 2. prefetch — chunk+1 from global into temp (in flight during math) + // 3. MMA math — on slice (chunk%2) + // 4. store — temp (chunk+1, dequantized) into slice ((chunk+1)%2) + // ========================================================= + uint chunk; + for (chunk = 0; chunk + 1u < num_chunks; ++chunk) { + const uint cur_base_A = (chunk % 2u) * ASH_SLICE; + const uint cur_base_B = (chunk % 2u) * BSH_SLICE; + const uint nxt_base_A = ((chunk + 1u) % 2u) * ASH_SLICE; + const uint nxt_base_B = ((chunk + 1u) % 2u) * BSH_SLICE; + + barrier(); + + // --- prefetch chunk+1 -> temp --- + { + const uint chunkK_nxt = (chunk + 1u) * WG_TILE_K; + + [[unroll]] for (uint p = 0; p < A_PASSES; ++p) { + const uint row = tile_m_start + p * A_ROWS_PER_PASS + a_row_offset; + const uint k_hv4 = (chunkK_nxt + a_col * FP16_PER_VEC4) / 4u; + temp_A[p] = load_a_vec4(row, k_hv4, K4); + } + [[unroll]] for (uint p = 0; p < B_PASSES; ++p) { + const uint k_row = chunkK_nxt + p * B_ROWS_PER_PASS + b_row_offset; +#ifdef WEIGHT_BUFFER + temp_B[p] = t_packed_weight[n8_blk * K4 + (k_row >> 2u)]; +#else + temp_B[p] = texelFetch(t_packed_weight, ivec2(k_row >> 2u, n8_blk), 0); +#endif + } + const uint group_nxt = (chunk + 1u) / CHUNKS_PER_GROUP; + if (group_nxt != cached_group) { + cached_group = group_nxt; + sc0 = t_weight_scales[group_nxt * N4 + sc_n4]; + sc1 = t_weight_scales[group_nxt * N4 + sc_n4 + 1u]; + } + } + + // --- MMA math on the cur slice --- + [[unroll]] for (uint k = 0; k < WG_TILE_K / MMA_K; ++k) { + const uint k_start = MMA_K * k; + + coopmat matA[MMAS_PER_SG_M]; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + const uint row_a = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + coopMatLoad( + matA[i], Ash, + cur_base_A + row_a * A_STRIDE_VEC4 + k_start / FP16_PER_VEC4, + A_STRIDE_VEC4, + gl_CooperativeMatrixLayoutRowMajor); + } + + coopmat matB; + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint col_b = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j) / FP16_PER_VEC4; + coopMatLoad( + matB, Bsh, + cur_base_B + k_start * B_STRIDE_VEC4 + col_b, + B_STRIDE_VEC4, + gl_CooperativeMatrixLayoutRowMajor); + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + result[i][j] = coopMatMulAdd(matA[i], matB, result[i][j]); + } + } + } + + // --- store temp (chunk+1) -> nxt slice, dequantizing B --- + { + [[unroll]] for (uint p = 0; p < A_PASSES; ++p) { + Ash[nxt_base_A + (p * A_ROWS_PER_PASS + a_row_offset) * A_STRIDE_VEC4 + a_col] = + temp_A[p]; + } + [[unroll]] for (uint p = 0; p < B_PASSES; ++p) { + Bsh[nxt_base_B + (p * B_ROWS_PER_PASS + b_row_offset) * B_STRIDE_VEC4 + b_col] = + dequant_block(temp_B[p], col_lo, col_hi, sc0, sc1); + } + } + } + + // --- epilogue: barrier, then MMA on the last chunk (loop peeled) --- + { + const uint cur_base_A = (chunk % 2u) * ASH_SLICE; + const uint cur_base_B = (chunk % 2u) * BSH_SLICE; + + barrier(); + + [[unroll]] for (uint k = 0; k < WG_TILE_K / MMA_K; ++k) { + const uint k_start = MMA_K * k; + + coopmat matA[MMAS_PER_SG_M]; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + const uint row_a = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + coopMatLoad( + matA[i], Ash, + cur_base_A + row_a * A_STRIDE_VEC4 + k_start / FP16_PER_VEC4, + A_STRIDE_VEC4, + gl_CooperativeMatrixLayoutRowMajor); + } + + coopmat matB; + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint col_b = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j) / FP16_PER_VEC4; + coopMatLoad( + matB, Bsh, + cur_base_B + k_start * B_STRIDE_VEC4 + col_b, + B_STRIDE_VEC4, + gl_CooperativeMatrixLayoutRowMajor); + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + result[i][j] = coopMatMulAdd(matA[i], matB, result[i][j]); + } + } + } + } + + // --- Bias staging (if any) --- +#ifdef HAS_BIAS + if (apply_bias > 0) { + for (uint t = gl_LocalInvocationID.x; t < WG_TILE_N; t += WG_SIZE) { + bias_sh[t] = float16_t(t_bias[tile_n_start + t]); + } + memoryBarrierShared(); + barrier(); + } +#endif + + // --- Store result tile --- +#ifdef IO_TEXTURE + // Epilogue iteration i drains accumulator row-block i from EVERY subgroup + // into Csh at once, so the SG_GRID_Y bands it holds are disjoint global row + // ranges; the whole workgroup then imageStores them. The band-to-global-row + // map below reproduces the buffer path's gi exactly: + // buffer: gi = tile_m_start + MMA_M * (MMAS_PER_SG_M * warpInTile.y + i) + // = tile_m_start + warpInTile.y * SG_TILE_M + i * MMA_M + // and lr / MMA_M is the writing subgroup's warpInTile.y. + // + // PORTABILITY NOTE: [[unroll]] is only a hint and glslc does not honor it + // here -- the barrier() in the loop body keeps the loop rolled, so + // result[i][j] IS dynamically indexed. Coopmat arrays are opaque per-lane + // storage and dynamic indexing is exactly the construct the Xclipse/AMD-PAL + // compiler has broken before, so check this first if the texture variants + // miscompile on M51. Fully unrolling would need the drain hand-expanded so + // each i gets its own barrier. + const uint CSH_TEXELS_PER_ROW = WG_TILE_N / 4u; + const uint CSH_TEXELS = CSH_ROWS * CSH_TEXELS_PER_ROW; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + // Guards Csh against the previous iteration's readers. Inert on i == 0 but + // must stay unconditional to remain workgroup-uniform. + barrier(); + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { +#ifdef HAS_BIAS + if (apply_bias > 0) { + const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopmat bias_tile; + coopMatLoad( + bias_tile, bias_sh, + local_n, /*stride=*/0u, + gl_CooperativeMatrixLayoutRowMajor); + result[i][j] += bias_tile; + } +#endif + coopMatStore( + result[i][j], Csh, + warpInTile.y * MMA_M * WG_TILE_N + + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j), + WG_TILE_N, + gl_CooperativeMatrixLayoutRowMajor); + } + memoryBarrierShared(); + barrier(); + + for (uint t = gl_LocalInvocationID.x; t < CSH_TEXELS; t += WG_SIZE) { + const uint lr = t / CSH_TEXELS_PER_ROW; + const uint lc4 = t % CSH_TEXELS_PER_ROW; + const uint m = + tile_m_start + (lr / MMA_M) * SG_TILE_M + i * MMA_M + (lr % MMA_M); + const uint base = lr * WG_TILE_N + lc4 * 4u; + imageStore( + t_output, + ivec3(tile_n_start / 4u + lc4, m, 0), + vec4( + float(Csh[base]), + float(Csh[base + 1u]), + float(Csh[base + 2u]), + float(Csh[base + 3u]))); + } + } +#else + const uint N_out = uint(out_N_arg); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint gi = tile_m_start + MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + const uint gj = tile_n_start + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + +#ifdef HAS_BIAS + if (apply_bias > 0) { + const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopmat bias_tile; + coopMatLoad( + bias_tile, bias_sh, + local_n, /*stride=*/0u, + gl_CooperativeMatrixLayoutRowMajor); + result[i][j] += bias_tile; + } +#endif + + coopmat out_tile = + coopmat(result[i][j]); + coopMatStore( + out_tile, t_output, + gi * N_out + gj, N_out, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#endif // IO_TEXTURE +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_coopmat_tsweep_dbuf4.yaml b/backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_coopmat_tsweep_dbuf4.yaml new file mode 100644 index 00000000000..fbc01bc546f --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_coopmat_tsweep_dbuf4.yaml @@ -0,0 +1,858 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# specs/041-dbuf4-tile-sweep: TILE-SIZE SWEEP variants of the fp16 q4gsw +# coopmat kernel's dbuf4 ("store-first", prefetch-peeled) loop structure -- +# linear_q4gsw_coopmat_tsweep_dbuf4.glsl is a fork of +# linear_q4gsw_coopmat_tsweep.glsl (dbuf1) with only the loop structure +# swapped. Only the tile geometry (WG_TILE_*, SG_GRID_*, SUBGROUP_SIZE) +# varies per variant. Selected at dispatch via +# ET_VK_Q4GSW_COOPMAT_VARIANT=tsweep_dbuf4_txkgs +# (QuantizedLinear.cpp). Seed variant below matches the current production +# 4w tile (dbuf1, specs/036) as a legal, known-fast starting point for +# sweep.py's Optuna search; specs/041's sweep appends further candidates +# here as it runs. + +linear_q4gsw_coopmat_tsweep_dbuf4: + parameter_names_with_default_values: + PRECISION: highp + HAS_BIAS: false + IO_STORAGE: buffer + WEIGHT_STORAGE: texture2d + MMA_M: 16 + MMA_N: 16 + MMA_K: 16 + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + shader_variants: + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t128x128k16g22s32_buffer_texture2d_half + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t128x128k16g22s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t128x128k16g22s32_buffer_buffer_half + WEIGHT_STORAGE: buffer + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x128k16g41s32_buffer_texture2d_half + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 4 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x128k16g41s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 4 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t256x256k16g14s32_buffer_texture2d_half + WEIGHT_STORAGE: texture2d + WG_TILE_M: 256 + WG_TILE_N: 256 + WG_TILE_K: 16 + SG_GRID_X: 1 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t256x256k16g14s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + WG_TILE_M: 256 + WG_TILE_N: 256 + WG_TILE_K: 16 + SG_GRID_X: 1 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x256k32g81s32_buffer_texture2d_half + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 256 + WG_TILE_K: 32 + SG_GRID_X: 8 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x256k32g81s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 256 + WG_TILE_K: 32 + SG_GRID_X: 8 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t32x64k64g41s64_buffer_texture2d_half + WEIGHT_STORAGE: texture2d + WG_TILE_M: 32 + WG_TILE_N: 64 + WG_TILE_K: 64 + SG_GRID_X: 4 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 64 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t32x64k64g41s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + WG_TILE_M: 32 + WG_TILE_N: 64 + WG_TILE_K: 64 + SG_GRID_X: 4 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 64 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x128k64g14s32_buffer_texture2d_half + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 1 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x128k64g14s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 1 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t128x256k32g18s32_buffer_texture2d_half + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 256 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 8 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t128x256k32g18s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 256 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 8 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t16x64k64g11s32_buffer_texture2d_half + WEIGHT_STORAGE: texture2d + WG_TILE_M: 16 + WG_TILE_N: 64 + WG_TILE_K: 64 + SG_GRID_X: 1 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t16x64k64g11s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + WG_TILE_M: 16 + WG_TILE_N: 64 + WG_TILE_K: 64 + SG_GRID_X: 1 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t256x128k32g44s32_buffer_texture2d_half + WEIGHT_STORAGE: texture2d + WG_TILE_M: 256 + WG_TILE_N: 128 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t256x128k32g44s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + WG_TILE_M: 256 + WG_TILE_N: 128 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x128k16g22s32_buffer_texture2d_half + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x128k16g22s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t32x64k128g12s32_buffer_texture2d_half + WEIGHT_STORAGE: texture2d + WG_TILE_M: 32 + WG_TILE_N: 64 + WG_TILE_K: 128 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t32x64k128g12s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + WG_TILE_M: 32 + WG_TILE_N: 64 + WG_TILE_K: 128 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x256k16g12s32_buffer_texture2d_half + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 256 + WG_TILE_K: 16 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x256k16g12s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 256 + WG_TILE_K: 16 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t16x64k64g21s64_buffer_texture2d_half + WEIGHT_STORAGE: texture2d + WG_TILE_M: 16 + WG_TILE_N: 64 + WG_TILE_K: 64 + SG_GRID_X: 2 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 64 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t16x64k64g21s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + WG_TILE_M: 16 + WG_TILE_N: 64 + WG_TILE_K: 64 + SG_GRID_X: 2 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 64 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x32k64g22s64_buffer_texture2d_half + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 32 + WG_TILE_K: 64 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x32k64g22s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 32 + WG_TILE_K: 64 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t128x256k16g81s32_buffer_texture2d_half + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 256 + WG_TILE_K: 16 + SG_GRID_X: 8 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t128x256k16g81s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 256 + WG_TILE_K: 16 + SG_GRID_X: 8 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t16x32k128g21s64_buffer_texture2d_half + WEIGHT_STORAGE: texture2d + WG_TILE_M: 16 + WG_TILE_N: 32 + WG_TILE_K: 128 + SG_GRID_X: 2 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 64 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t16x32k128g21s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + WG_TILE_M: 16 + WG_TILE_N: 32 + WG_TILE_K: 128 + SG_GRID_X: 2 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 64 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x128k64g22s64_buffer_texture2d_half + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x128k64g22s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x128k64g81s32_buffer_texture2d_half + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 8 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x128k64g81s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 8 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t32x128k64g41s32_buffer_texture2d_half + WEIGHT_STORAGE: texture2d + WG_TILE_M: 32 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 4 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t32x128k64g41s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + WG_TILE_M: 32 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 4 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t16x64k32g11s64_buffer_texture2d_half + WEIGHT_STORAGE: texture2d + WG_TILE_M: 16 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 64 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t16x64k32g11s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + WG_TILE_M: 16 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 64 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t16x16k64g11s64_buffer_texture2d_half + WEIGHT_STORAGE: texture2d + WG_TILE_M: 16 + WG_TILE_N: 16 + WG_TILE_K: 64 + SG_GRID_X: 1 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 64 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t16x16k64g11s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + WG_TILE_M: 16 + WG_TILE_N: 16 + WG_TILE_K: 64 + SG_GRID_X: 1 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 64 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t128x64k16g22s32_buffer_texture2d_half + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 16 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t128x64k16g22s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 16 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t32x128k32g22s32_buffer_texture2d_half + WEIGHT_STORAGE: texture2d + WG_TILE_M: 32 + WG_TILE_N: 128 + WG_TILE_K: 32 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t32x128k32g22s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + WG_TILE_M: 32 + WG_TILE_N: 128 + WG_TILE_K: 32 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x128k32g11s32_buffer_texture2d_half + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x128k32g11s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t128x64k64g12s32_buffer_texture2d_half + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 64 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t128x64k64g12s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 64 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x128k16g12s64_buffer_texture2d_half + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x128k16g12s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t128x128k16g21s32_buffer_texture2d_half + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 2 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t128x128k16g21s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 2 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t16x64k32g21s32_buffer_texture2d_half + WEIGHT_STORAGE: texture2d + WG_TILE_M: 16 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 2 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t16x64k32g21s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + WG_TILE_M: 16 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 2 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x128k16g14s32_buffer_texture2d_half + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 1 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x128k16g14s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 1 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t16x32k64g11s32_buffer_texture2d_half + WEIGHT_STORAGE: texture2d + WG_TILE_M: 16 + WG_TILE_N: 32 + WG_TILE_K: 64 + SG_GRID_X: 1 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t16x32k64g11s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + WG_TILE_M: 16 + WG_TILE_N: 32 + WG_TILE_K: 64 + SG_GRID_X: 1 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x32k64g21s64_buffer_texture2d_half + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 32 + WG_TILE_K: 64 + SG_GRID_X: 2 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 64 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x32k64g21s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 32 + WG_TILE_K: 64 + SG_GRID_X: 2 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 64 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t32x64k64g11s32_buffer_texture2d_half + WEIGHT_STORAGE: texture2d + WG_TILE_M: 32 + WG_TILE_N: 64 + WG_TILE_K: 64 + SG_GRID_X: 1 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t32x64k64g11s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + WG_TILE_M: 32 + WG_TILE_N: 64 + WG_TILE_K: 64 + SG_GRID_X: 1 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t128x64k16g21s32_buffer_texture2d_half + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 16 + SG_GRID_X: 2 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t128x64k16g21s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 16 + SG_GRID_X: 2 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x128k64g21s64_buffer_texture2d_half + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 2 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 64 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x128k64g21s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 2 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 64 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x128k32g22s64_buffer_texture2d_half + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 32 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x128k32g22s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 32 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t32x128k64g22s64_buffer_texture2d_half + WEIGHT_STORAGE: texture2d + WG_TILE_M: 32 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t32x128k64g22s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + WG_TILE_M: 32 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t128x128k16g14s32_buffer_texture2d_half + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 1 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t128x128k16g14s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 1 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t128x64k16g12s32_buffer_texture2d_half + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 16 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t128x64k16g12s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 16 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x128k16g21s32_buffer_texture2d_half + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 2 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x128k16g21s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 2 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t128x64k32g24s32_buffer_texture2d_half + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t128x64k32g24s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t32x64k32g22s32_buffer_texture2d_half + WEIGHT_STORAGE: texture2d + WG_TILE_M: 32 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t32x64k32g22s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + WG_TILE_M: 32 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t32x256k32g21s32_buffer_texture2d_half + WEIGHT_STORAGE: texture2d + WG_TILE_M: 32 + WG_TILE_N: 256 + WG_TILE_K: 32 + SG_GRID_X: 2 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t32x256k32g21s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + WG_TILE_M: 32 + WG_TILE_N: 256 + WG_TILE_K: 32 + SG_GRID_X: 2 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x128k64g41s32_buffer_texture2d_half + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 4 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x128k64g41s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 4 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t128x32k64g22s32_buffer_texture2d_half + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 32 + WG_TILE_K: 64 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t128x32k64g22s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 32 + WG_TILE_K: 64 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t16x64k128g21s32_buffer_texture2d_half + WEIGHT_STORAGE: texture2d + WG_TILE_M: 16 + WG_TILE_N: 64 + WG_TILE_K: 128 + SG_GRID_X: 2 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t16x64k128g21s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + WG_TILE_M: 16 + WG_TILE_N: 64 + WG_TILE_K: 128 + SG_GRID_X: 2 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t128x64k64g42s32_buffer_texture2d_half + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 64 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t128x64k64g42s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 64 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t128x128k32g42s32_buffer_texture2d_half + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t128x128k32g42s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x64k64g11s64_buffer_texture2d_half + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 64 + WG_TILE_K: 64 + SG_GRID_X: 1 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 64 + - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x64k64g11s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 64 + WG_TILE_K: 64 + SG_GRID_X: 1 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 64 diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_qw_coopmat.yaml b/backends/vulkan/runtime/graph/ops/glsl/linear_qw_coopmat.yaml index 173caf8dc29..cd17e9edd6c 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_qw_coopmat.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_qw_coopmat.yaml @@ -10,11 +10,18 @@ # Forces buffer storage for activation/output (coopMatLoad/Store on buffers); # INT weight storage can be texture2d or buffer (matches the tiled path). # DTYPE = half only; fp32 activations are not supported. -# Geometry follows the double-buffered reference (coopmat_mm_ref): 128x64 -# tile, K-step 16, 2x2 subgroup grid x 32 threads (subgroup size 32 forced). -# This geometry was tuned for AMD-RDNA GPUs. NOTE: the C++ dispatch in +# Geometry follows the double-buffered reference (coopmat_mm_ref): 128x128 +# tile, K-step 16, 4 subgroups x 32 threads (subgroup size 32 forced). The +# 128x128 / 2x2-subgroup-grid geometry is specs/036-portable-device-sweep's +# e2e-ranked winner on M51 (samsung xclipse 970), re-verified through this +# shipped shader path (not just the tsweep_ toggle) 2026-07-24 — confirmed +# +6.8%/+7.7%/+10.1% (1B/3B/8B prefill tok/s) over the prior 128x64/2x2 tile +# (kept as tsweep seed tsweep_t128x64k16g22s32 for future sweeps). NOTE: an +# EARLIER 128x128 shape (4x2 subgroup grid, not 2x2) was -25% vs 128x64/2x2 +# on the same device — don't conflate the two; the subgroup grid, not just +# the tile size, is what makes this one different. NOTE: the C++ dispatch in # QuantizedLinear.cpp must keep kQ4gswCoopmatDims.n and .wg_size in sync with -# WG_TILE_N (64) and WG_SIZE (= SG_GRID_X*SG_GRID_Y*SUBGROUP = 128). +# WG_TILE_N (128) and WG_SIZE (= SG_GRID_X*SG_GRID_Y*SUBGROUP = 128). linear_qw_coopmat: parameter_names_with_default_values: @@ -26,7 +33,7 @@ linear_qw_coopmat: MMA_N: 16 MMA_K: 16 WG_TILE_M: 128 - WG_TILE_N: 64 + WG_TILE_N: 128 WG_TILE_K: 16 SG_GRID_X: 2 SG_GRID_Y: 2 diff --git a/backends/vulkan/runtime/graph/ops/impl/Q4gswLinear.cpp b/backends/vulkan/runtime/graph/ops/impl/Q4gswLinear.cpp index 62322602ac3..108cef02c06 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Q4gswLinear.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Q4gswLinear.cpp @@ -676,7 +676,13 @@ void q4gsw_linear(ComputeGraph& graph, const std::vector& args) { REGISTER_OPERATORS { VK_REGISTER_OP(et_vk.q4gsw_linear.default, q4gsw_linear); - VK_REGISTER_OP(et_vk.linear_q4gsw.default, q4gsw_linear); + // et_vk.linear_q4gsw.default is registered by QuantizedLinear.cpp instead + // (Option B, memory quant-perf-rebase-orphaned-4w-coopmat) -- this file's + // q4gsw_linear_gemm__* shaders are tiled-only with no coopmat path, and + // registering here would silently make QuantizedLinear.cpp's + // linear_q4gsw_coopmat (present, correctly built) unreachable for every 4w + // PTE. + // VK_REGISTER_OP(et_vk.linear_q4gsw.default, q4gsw_linear); } } // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp index 45588e7e2e5..75d4eb6abcb 100644 --- a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp @@ -8,6 +8,8 @@ #include +#include + #include #include #include @@ -53,10 +55,15 @@ void resize_linear_qw_node( // Per-shader coopmat tile geometry (must match each shader's yaml). // Workgroup size (wg_size) = SG_GRID_X * SG_GRID_Y * SUBGROUP_SIZE. -// linear_q4gsw_coopmat 128x64x16, 2x2 subgroups x 32 (forced) -> 128 -// linear_dq8ca_q4gsw_coopmat 128x64x32, 2x2 subgroups x 64 -> 256 -// (The int8-MMA shaders stay on wave64: int8 WMMA at forced subgroup 32 -// crashes the Xclipse PAL compiler.) +// linear_q4gsw_coopmat 128x128x16, 2x2 subgroups x 32 (forced) -> 128 +// linear_dq8ca_q4gsw_coopmat 64x32x32, 1x2 subgroups x 64 -> 128 +// (specs/027-e2e-tile-sweep: dq8ca_q4gsw tile updated from 128x64x32/2x2 to +// this e2e-ranked winner. int8-MMA stays on wave64 at this tile -- specs/026 +// found subgroup=32 legal-but-shape-dependently-incorrect, not verified +// correct at this specific tile. specs/036-portable-device-sweep: q4gsw tile +// updated 2026-07-24 from 128x64x16/2x2 to this e2e-ranked winner, +6.8/+7.7/ +// +10.1% on 1B/3B/8B prefill tok/s -- N doubled, grid/subgroup unchanged so +// wg_size is unaffected.) struct CoopmatTileDims { uint32_t m; uint32_t n; @@ -65,20 +72,113 @@ struct CoopmatTileDims { // the WG_SIZE the shader yaml resolves to, or the launched thread count won't // match the shader's staging passes (out-of-bounds). uint32_t wg_size; + // Only needed for the texture-IO shared-memory budget below; the buffer path + // never reads it. 0 = "unknown / shipped default". + uint32_t sg_grid_y; +}; +// linear_qw_coopmat.yaml: 128x128, 2x2 subgroup grid, sg32 -> WG_SIZE 128. +constexpr CoopmatTileDims kQ4gswCoopmatDims = {128, 128, 16, 128, 2}; +// linear_dq8ca_qw_coopmat.yaml: 64x32, 1x2 grid, sg64 -> WG_SIZE 128 +// (specs/027-e2e-tile-sweep winner, was 128x64x32/256). +constexpr CoopmatTileDims kDq8caQ4gswCoopmatDims = {64, 32, 32, 128, 2}; + +// specs/028-4w-e2e-tile-sweep / specs/041-dbuf4-tile-sweep: +// ET_VK_Q4GSW_COOPMAT_VARIANT / ET_VK_DQ8CA_COOPMAT_VARIANT can swap the +// coopmat dispatch to a specific tile/subgroup-grid/loop-structure variant for +// sweeping. A "tsweep_dbuf_t..." token additionally selects a loop-structure +// variant (1-4); "tsweep_t..." is the original (production winner's own) +// namespace. Unset/unrecognized = shipped dispatch, unchanged. The five +// prefixes are mutually exclusive by construction (position 7 is 'd' vs 't'). +static const char* const kTsweepPrefixes[] = { + "tsweep_dbuf1_t", + "tsweep_dbuf2_t", + "tsweep_dbuf3_t", + "tsweep_dbuf4_t", + "tsweep_t", }; -// linear_qw_coopmat.yaml: 128x64, 2x2 subgroup grid, sg32 -> WG_SIZE 128. -constexpr CoopmatTileDims kQ4gswCoopmatDims = {128, 64, 16, 128}; -// linear_dq8ca_qw_coopmat.yaml: 128x64, 2x2 grid, sg64 -> WG_SIZE 256. -constexpr CoopmatTileDims kDq8caQ4gswCoopmatDims = {128, 64, 32, 256}; + +static bool is_recognized_coopmat_variant_token(const std::string& v) { + for (const char* prefix : kTsweepPrefixes) { + if (v.rfind(prefix, 0) == 0) { + return true; + } + } + return false; +} + +static const std::string& q4gsw_coopmat_variant() { + static const std::string variant = [] { + const char* env = std::getenv("ET_VK_Q4GSW_COOPMAT_VARIANT"); + if (!env) { + return std::string(); + } + const std::string v(env); + if (is_recognized_coopmat_variant_token(v)) { + return v; + } + return std::string(); + }(); + return variant; +} + +static const std::string& dq8ca_coopmat_variant() { + static const std::string variant = [] { + const char* env = std::getenv("ET_VK_DQ8CA_COOPMAT_VARIANT"); + if (!env) { + return std::string(); + } + const std::string v(env); + if (is_recognized_coopmat_variant_token(v)) { + return v; + } + return std::string(); + }(); + return variant; +} + +// Parses "tsweep_txkgs" or +// "tsweep_dbuf_txkgs" -> {M, N, K, SGX*SGY*sub, +// SGY}. Returns fallback unchanged if the token matches none of +// kTsweepPrefixes. +static CoopmatTileDims parse_tsweep_tile( + const std::string& variant, + const CoopmatTileDims& fallback) { + size_t t_pos = std::string::npos; + for (const char* prefix : kTsweepPrefixes) { + if (variant.rfind(prefix, 0) == 0) { + t_pos = std::strlen(prefix); + break; + } + } + if (t_pos == std::string::npos) { + return fallback; + } + const size_t x_pos = variant.find('x', t_pos); + const size_t k_pos = variant.find('k', x_pos); + const size_t g_pos = variant.find('g', k_pos); + const size_t s_pos = variant.find('s', g_pos); + const uint32_t m = std::stoul(variant.substr(t_pos, x_pos - t_pos)); + const uint32_t n = std::stoul(variant.substr(x_pos + 1, k_pos - x_pos - 1)); + const uint32_t k = std::stoul(variant.substr(k_pos + 1, g_pos - k_pos - 1)); + const std::string grid = variant.substr(g_pos + 1, s_pos - g_pos - 1); + const uint32_t sgx = grid[0] - '0'; + const uint32_t sgy = grid[1] - '0'; + const uint32_t sub = std::stoul(variant.substr(s_pos + 1)); + return {m, n, k, sgx * sgy * sub, sgy}; +} + +static CoopmatTileDims parse_q4gsw_tsweep_tile(const std::string& variant) { + return parse_tsweep_tile(variant, kQ4gswCoopmatDims); +} static CoopmatTileDims coopmat_tile_dims(const std::string& kernel_name) { // Exact prefix matches (the "linear_dq8ca_*" names must not match the - // weight-only entries). - if (kernel_name.rfind("linear_q4gsw_coopmat", 0) == 0) { - return kQ4gswCoopmatDims; - } + // weight-only entries). Order matters: check dq8ca first. if (kernel_name.rfind("linear_dq8ca_q4gsw_coopmat", 0) == 0) { - return kDq8caQ4gswCoopmatDims; + return parse_tsweep_tile(dq8ca_coopmat_variant(), kDq8caQ4gswCoopmatDims); + } + if (kernel_name.rfind("linear_q4gsw_coopmat", 0) == 0) { + return parse_q4gsw_tsweep_tile(q4gsw_coopmat_variant()); } return {kCoopmatTileM, kCoopmatTileN, kCoopmatTileK, kCoopmatInvocations}; } @@ -152,6 +252,15 @@ utils::uvec3 quantized_linear_local_wg_size( } } +// Experiment hook (specs/040/041): allows the *_texture3d_* variants, which +// stage the result tile through shared memory and imageStore it instead of +// coopMatStore-ing straight to an SSBO. Off by default, so buffer dispatch is +// byte-identical. +static bool texture_coopmat_enabled() { + static const bool enabled = std::getenv("ET_VK_TEXTURE_COOPMAT") != nullptr; + return enabled; +} + // Returns true when the q4gsw coopmat shader can be dispatched for this // (M, N, K, dtype, output_storage, group_size) tuple. Preconditions match what // linear_q4gsw_coopmat.glsl assumes; the subgroup_size == 64 check scopes this @@ -164,7 +273,15 @@ static bool can_use_q4gsw_coopmat( const ValueRef bias, int64_t tile_m = kCoopmatTileM, int64_t tile_n = kCoopmatTileN, - int64_t tile_k = kCoopmatTileK) { + int64_t tile_k = kCoopmatTileK, + bool allow_texture_io = false, + uint32_t sg_grid_y = 0) { + // Baseline-measurement escape hatch: forces every dispatch through this + // function to the tiled fallback, regardless of eligibility. Off by + // default (unset), so production behavior is unchanged. + if (std::getenv("ET_VK_FORCE_TILED_LINEAR") != nullptr) { + return false; + } // The coopmat shaders only build HAS_BIAS=false variants, so they would // silently drop a bias. Fall back to the tiled path (which applies bias at // runtime via the apply_bias spec constant) whenever a bias is present. @@ -185,19 +302,60 @@ static bool can_use_q4gsw_coopmat( if (!graph->device_is_amd()) { return false; } - // Coopmat shaders dispatch over gl_WorkGroupID.xy only; batched (rank > 2) - // outputs would silently miscompute all slices beyond the first. - if (graph->dim_of(output) > 2) { + // Coopmat shaders dispatch over gl_WorkGroupID.xy only, sized purely from + // the output's trailing two dims; neither that sizing nor the shaders + // themselves ever read a leading dim. A genuine batch (any leading-dim + // product != 1) would silently miscompute all slices beyond the first -- + // but a size-1 leading dim (the real exported model's rank-3 [1, M, K] + // activations, never squeezed) is safe: a contiguous Buffer's [1, M, N] + // layout is bit-identical to [M, N] when the leading dim is 1, so the + // existing 2D dispatch grid already covers 100% of the data. Reject only a + // real batch. + const std::vector out_sizes = graph->sizes_of(output); + int64_t leading_dims_numel = 1; + for (int64_t d = 0; d < graph->dim_of(output) - 2; d++) { + leading_dims_numel *= utils::val_at(d, out_sizes); + } + if (leading_dims_numel != 1) { return false; } if (graph->storage_type_of(output) != utils::kBuffer) { - return false; + // One IO_STORAGE param is shared across t_input/t_output, so BOTH must be + // texture3d; the imageStore epilogue and texelFetch A-stage assume + // width-packed texels. + if (!allow_texture_io || !texture_coopmat_enabled()) { + return false; + } + if (graph->storage_type_of(output) != utils::kTexture3D || + graph->storage_type_of(fp_input) != utils::kTexture3D) { + return false; + } + if (graph->packed_dim_of(output) != WHCN::kWidthDim || + graph->packed_dim_of(fp_input) != WHCN::kWidthDim) { + return false; + } + // The texture epilogue needs a Csh staging array ON TOP OF the Ash/Bsh the + // buffer path already allocates: SG_GRID_Y * MMA_M rows x WG_TILE_N fp16. + // That term is absent from the offline tile_constraints model, so a tile + // that is legal for buffer can exceed the shared-memory limit at texture + // IO -- one specific large tile hung the GPU and rebooted the board + // instead of failing pipeline creation (2026-08-09). Reject here so it + // falls back to tiled. + if (sg_grid_y > 0) { + constexpr int64_t kMmaM = 16; // MMA_M, fixed across every coopmat yaml + const int64_t csh_bytes = + int64_t(sg_grid_y) * kMmaM * tile_n * int64_t(sizeof(uint16_t)); + const int64_t limit = + graph->context()->adapter_ptr()->max_compute_shared_memory_size(); + if (csh_bytes >= limit) { + return false; + } + } } if (graph->dtype_of(output) != vkapi::kHalf) { return false; } - const std::vector out_sizes = graph->sizes_of(output); const int64_t N = utils::val_at(-1, out_sizes); const int64_t M = utils::val_at(-2, out_sizes); const std::vector in_sizes = graph->sizes_of(fp_input); @@ -236,18 +394,29 @@ vkapi::ShaderInfo pick_linear_qw_shader( if (weight_is_4bit && !is_gemv_case) { const int64_t group_size = graph->extract_scalar(resize_args.at(0)); + // A tsweep_* variant has different tile dims than the shipped + // kQ4gswCoopmatDims, so the eligibility check's alignment gate must use + // the ACTIVE variant's own dims, not the shipped constant. + const CoopmatTileDims active_dims = + parse_q4gsw_tsweep_tile(q4gsw_coopmat_variant()); if (can_use_q4gsw_coopmat( graph, output, fp_input, group_size, resize_args.at(2), - kQ4gswCoopmatDims.m, - kQ4gswCoopmatDims.n, - kQ4gswCoopmatDims.k)) { + active_dims.m, + active_dims.n, + active_dims.k, + /*allow_texture_io=*/true, + active_dims.sg_grid_y)) { std::string kernel_name = "linear_q4gsw_coopmat"; - // Output storage is buffer (gated above); weight storage matches the - // existing variants. + const std::string& variant = q4gsw_coopmat_variant(); + if (!variant.empty()) { + kernel_name += "_" + variant; + } + // Output storage is buffer or texture3d (gated above); weight storage + // matches the existing variants. add_storage_type_suffix(kernel_name, graph->storage_type_of(output)); add_storage_type_suffix( kernel_name, graph->storage_type_of(packed_int_weight)); @@ -299,16 +468,33 @@ vkapi::ShaderInfo pick_linear_dqa_qw_shader( graph->context()->adapter_ptr()->supports_int8_cooperative_matrix()) { const int64_t group_size = graph->extract_scalar(resize_args.at(0)); + // Alignment gate must use the ACTIVE sweep variant's own tile dims (same + // rationale as the q4gsw tsweep hook above). + const CoopmatTileDims active_dq8ca_dims = + parse_tsweep_tile(dq8ca_coopmat_variant(), kDq8caQ4gswCoopmatDims); + // The dq8ca texture-IO shader declares t_input (fp_input) with an + // IO_STORAGE-typed binding even though it's never read in the shader body + // (activations arrive pre-quantized in t_packed_int8_input instead) -- + // Vulkan still requires the bound resource's storage type to match the + // declared binding type, so fp_input must genuinely be texture3d too when + // texture IO is active. Same requirement as q4gsw; no separate check + // needed. if (can_use_q4gsw_coopmat( graph, out, fp_input, group_size, resize_args.at(2), - kDq8caQ4gswCoopmatDims.m, - kDq8caQ4gswCoopmatDims.n, - kDq8caQ4gswCoopmatDims.k)) { + active_dq8ca_dims.m, + active_dq8ca_dims.n, + active_dq8ca_dims.k, + /*allow_texture_io=*/true, + active_dq8ca_dims.sg_grid_y)) { std::string kernel_name = "linear_dq8ca_q4gsw_coopmat"; + const std::string& dq8ca_variant = dq8ca_coopmat_variant(); + if (!dq8ca_variant.empty()) { + kernel_name += "_" + dq8ca_variant; + } add_storage_type_suffix(kernel_name, graph->storage_type_of(out)); add_storage_type_suffix(kernel_name, graph->storage_type_of(int_weight)); add_dtype_suffix(kernel_name, graph->dtype_of(out)); @@ -940,6 +1126,44 @@ void linear_q8csw(ComputeGraph& graph, const std::vector& args) { output); } +// Registered below as et_vk.linear_q4gsw.default -- takes over from +// Q4gswLinear.cpp's own registration of the same op name (commented out +// there) so that add_linear_qw_node / linear_q4gsw_coopmat are actually +// reachable. See memory quant-perf-rebase-orphaned-4w-coopmat: upstream's +// Q4gswLinear.cpp silently hijacks this exact op name with its own +// q4gsw_linear_gemm__* tiled-only shaders, and QuantizedLinear.cpp's coopmat +// path (present and correctly built) is unreachable for any 4w PTE until +// this registration is restored. +void linear_q4gsw(ComputeGraph& graph, const std::vector& args) { + int32_t idx = 0; + const ValueRef fp_input = args.at(idx++); + const ValueRef weight_data = args.at(idx++); + const ValueRef weight_scales_data = args.at(idx++); + const ValueRef group_size = args.at(idx++); + const ValueRef bias_data = args.at(idx++); + const ValueRef output = args.at(idx++); + + const int64_t group_size_val = graph.extract_scalar(group_size); + + QuantizationConfig input_quant_config(32, kNoQuantization, {}); + QuantizationConfig weight_quant_config(4, kPerGroup, {group_size_val}); + + quantized_linear_impl( + graph, + input_quant_config, + weight_quant_config, + fp_input, + kDummyValueRef, // input scale + kDummyValueRef, // input zp + weight_data, + kDummyValueRef, // weight sums + weight_scales_data, + kDummyValueRef, // weight zeros + group_size, // group size + bias_data, + output); +} + void linear_dq8ca_q4gsw( ComputeGraph& graph, const std::vector& args) { @@ -978,6 +1202,7 @@ void linear_dq8ca_q4gsw( REGISTER_OPERATORS { VK_REGISTER_OP(et_vk.linear_q8ta_q8csw.default, linear_q8ta_q8csw); VK_REGISTER_OP(et_vk.linear_q8csw.default, linear_q8csw); + VK_REGISTER_OP(et_vk.linear_q4gsw.default, linear_q4gsw); VK_REGISTER_OP(et_vk.linear_dq8ca_q4gsw.default, linear_dq8ca_q4gsw); } diff --git a/backends/vulkan/runtime/vk_api/Adapter.h b/backends/vulkan/runtime/vk_api/Adapter.h index a1b7f2962ec..bd627aaea17 100644 --- a/backends/vulkan/runtime/vk_api/Adapter.h +++ b/backends/vulkan/runtime/vk_api/Adapter.h @@ -427,6 +427,10 @@ class Adapter final { return physical_device_.properties.limits.maxStorageBufferRange; } + inline uint32_t max_compute_shared_memory_size() const { + return physical_device_.properties.limits.maxComputeSharedMemorySize; + } + // Command Buffer Submission void submit_cmd( From 26f455784fbda8bf7c02b82d21f4af3380b385c3 Mon Sep 17 00:00:00 2001 From: Yanwen Xu Date: Tue, 18 Aug 2026 11:32:10 -0700 Subject: [PATCH 04/28] [ET-VK] Ship 8da4w texture-IO tile sweep winner as the new default Full precompiled-variant sweep (125 texture3d tiles, e2e prefill tok/s on real llama_main, not microbench) on M51 (secondary board, SUMD change-66942) found tsweep_dbuf4_t128x16k64g12s64 beats the prior default (t64x32k32g12s64) by 11.8-17.3% across 1B/3B/8B, with no rank flip. ETDump- confirmed genuine coopmat dispatch on the real prefill path both before and after this change. 4w's existing default (t128x128k16g22s32) was re-confirmed as still optimal in the same sweep -- no change there. Only takes effect when ET_VK_TEXTURE_COOPMAT=1 is set; that master switch stays opt-in. ET_VK_DQ8CA_COOPMAT_VARIANT still overrides this default when explicitly set. Authored with Claude. --- .../vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp index 75d4eb6abcb..7536adb84f2 100644 --- a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp @@ -122,16 +122,23 @@ static const std::string& q4gsw_coopmat_variant() { } static const std::string& dq8ca_coopmat_variant() { + // Default (no ET_VK_DQ8CA_COOPMAT_VARIANT set): + // tsweep_dbuf4_t128x16k64g12s64, this workspace's texture-IO tile sweep + // winner on M51 (2026-08-18) -- 11.8-17.3% faster than the prior + // t64x32k32g12s64 default across 1B/3B/8B, e2e-confirmed (not microbench) and + // ETDump-confirmed to actually dispatch coopmat on the real prefill path. + // Only takes effect when ET_VK_TEXTURE_COOPMAT=1 is also set -- that master + // switch stays opt-in. static const std::string variant = [] { const char* env = std::getenv("ET_VK_DQ8CA_COOPMAT_VARIANT"); if (!env) { - return std::string(); + return std::string("tsweep_dbuf4_t128x16k64g12s64"); } const std::string v(env); if (is_recognized_coopmat_variant_token(v)) { return v; } - return std::string(); + return std::string("tsweep_dbuf4_t128x16k64g12s64"); }(); return variant; } From 3ceeefc269dd2069f28ceed3051ee7a57788ad73 Mon Sep 17 00:00:00 2001 From: Yanwen Xu Date: Tue, 18 Aug 2026 11:39:01 -0700 Subject: [PATCH 05/28] [ET-VK] Fix q4gsw texture-IO default crash: give it a default variant too ET_VK_TEXTURE_COOPMAT=1 alone (no explicit ET_VK_Q4GSW_COOPMAT_VARIANT) crashed on every model size: the eligibility gate accepts texture3d storage unconditionally for q4gsw, but the empty-variant kernel name resolves to the bare "linear_q4gsw_coopmat" shader, which was never compiled with a texture3d IO_STORAGE variant (only the tsweep_dbuf4-suffixed one was) -- "Could not find ShaderInfo with name linear_q4gsw_coopmat_texture3d_texture2d_half". Same fix as dq8ca_coopmat_variant() (prior commit): default the empty case to tsweep_dbuf4_t128x128k16g22s32 -- same geometry as the shipped buffer default, already re-confirmed #1 in the M51 tile sweep, just resolving to a kernel name that actually has a texture3d build. Authored with Claude. --- .../runtime/graph/ops/impl/QuantizedLinear.cpp | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp index 7536adb84f2..cd8246c286f 100644 --- a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp @@ -107,16 +107,26 @@ static bool is_recognized_coopmat_variant_token(const std::string& v) { } static const std::string& q4gsw_coopmat_variant() { + // Default (no ET_VK_Q4GSW_COOPMAT_VARIANT set): + // tsweep_dbuf4_t128x128k16g22s32 + // -- same geometry as the shipped buffer-storage default (re-confirmed #1 in + // the same M51 tile sweep that picked dq8ca's new default), but the BARE + // "linear_q4gsw_coopmat" kernel name was never compiled with a texture3d + // IO_STORAGE variant (only the tsweep_dbuf4-suffixed shader was). Without + // this, ET_VK_TEXTURE_COOPMAT=1 alone (no explicit variant) resolves to a + // kernel name with no texture3d build and crashes at dispatch -- the gate + // accepts texture-IO but the bare name can't serve it. Same fix as + // dq8ca_coopmat_variant() below. static const std::string variant = [] { const char* env = std::getenv("ET_VK_Q4GSW_COOPMAT_VARIANT"); if (!env) { - return std::string(); + return std::string("tsweep_dbuf4_t128x128k16g22s32"); } const std::string v(env); if (is_recognized_coopmat_variant_token(v)) { return v; } - return std::string(); + return std::string("tsweep_dbuf4_t128x128k16g22s32"); }(); return variant; } From 0320ac3927923499f90385a0d9a7b93bf14df838 Mon Sep 17 00:00:00 2001 From: Yanwen Xu Date: Tue, 18 Aug 2026 12:09:04 -0700 Subject: [PATCH 06/28] [ET-VK] Revert 8da4w default to t64x32k32g12s64 -- new tile was numerically wrong The M51 tile-sweep winner shipped as the 8da4w default two commits ago (tsweep_dbuf4_t128x16k64g12s64) failed test_llama_microbench --correctness-only across every texture3d case, including the rank-3 (real model shape) ones -- roughly 70% of output elements mismatched in a structured per-row pattern. e2e prefill tok/s and ETDump dispatch confirmation are not a substitute for a real numeric correctness check; the sweep never ran one, and should have before this shipped. Checked and ruled out the M>K/2 int_input_sums undersizing described in upstream issue #21423: that buffer is marked unused in both the old and new tsweep_dbuf4 shaders identically, so it isn't the discriminator here. Root cause of the new tile's specific failure not yet isolated -- most likely a genuine indexing bug in the dbuf4 template's spec-resolved code at that exact tile shape (M=128, N=16, K=64, 1x2 grid, sub=64). Reverts to t64x32k32g12s64, now verified via test_llama_microbench --scheme=8da4w --storage=texture3d --correctness-only: 0 failures, real coopmat dispatch confirmed via the harness's own dispatch log (not just "produces coherent text"). Also ports test_llama_microbench.cpp onto this branch (from texture-dbuf4, where it originated) plus its CMakeLists registration, so this check can be re-run here going forward instead of depending on a cross-branch binary. Authored with Claude. --- .../graph/ops/impl/QuantizedLinear.cpp | 34 +- .../vulkan/test/custom_ops/CMakeLists.txt | 1 + .../test/custom_ops/test_llama_microbench.cpp | 1595 +++++++++++++++++ 3 files changed, 1621 insertions(+), 9 deletions(-) create mode 100644 backends/vulkan/test/custom_ops/test_llama_microbench.cpp diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp index cd8246c286f..68b9f89ea9b 100644 --- a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp @@ -132,23 +132,39 @@ static const std::string& q4gsw_coopmat_variant() { } static const std::string& dq8ca_coopmat_variant() { - // Default (no ET_VK_DQ8CA_COOPMAT_VARIANT set): - // tsweep_dbuf4_t128x16k64g12s64, this workspace's texture-IO tile sweep - // winner on M51 (2026-08-18) -- 11.8-17.3% faster than the prior - // t64x32k32g12s64 default across 1B/3B/8B, e2e-confirmed (not microbench) and - // ETDump-confirmed to actually dispatch coopmat on the real prefill path. - // Only takes effect when ET_VK_TEXTURE_COOPMAT=1 is also set -- that master - // switch stays opt-in. + // Default (no ET_VK_DQ8CA_COOPMAT_VARIANT set): tsweep_dbuf4_t64x32k32g12s64 + // -- same geometry as the shipped buffer-storage default, resolved through + // the tsweep_dbuf4 texture3d-capable shader. + // + // REVERTED 2026-08-18: the M51 tile sweep's apparent winner, + // tsweep_dbuf4_t128x16k64g12s64 (11.8-17.3% faster in e2e prefill tok/s), + // was shipped as this default for a short window and then found to be + // NUMERICALLY WRONG via test_llama_microbench --correctness-only -- + // every single texture3d correctness case failed, including the rank-3 + // (real model shape) ones (~70% of elements mismatched in a structured + // per-row pattern: the first M-rows correct, everything past a boundary + // wrong). Root cause not yet isolated -- checked and ruled out + // int_input_sums undersizing (github.com/pytorch/executorch/issues/21423): + // that buffer is marked unused in both this shader and the tsweep_dbuf4 + // one, identically, so it can't be the discriminator. Most likely a + // genuine indexing bug specific to that tile's spec-resolved shape + // (M=128, N=16, K=64, 1x2 subgroup grid, sub=64) in the dbuf4 template. + // This tile (t64x32k32g12s64) IS correctness-bench-confirmed clean: 0 + // failures across all texture3d cases including rank-3, real coopmat + // dispatch confirmed via the harness's own dispatch log. Do not re-ship + // t128x16k64g12s64 (or any other untested sweep candidate) without first + // passing test_llama_microbench --scheme=8da4w --storage=texture3d + // --correctness-only clean. static const std::string variant = [] { const char* env = std::getenv("ET_VK_DQ8CA_COOPMAT_VARIANT"); if (!env) { - return std::string("tsweep_dbuf4_t128x16k64g12s64"); + return std::string("tsweep_dbuf4_t64x32k32g12s64"); } const std::string v(env); if (is_recognized_coopmat_variant_token(v)) { return v; } - return std::string("tsweep_dbuf4_t128x16k64g12s64"); + return std::string("tsweep_dbuf4_t64x32k32g12s64"); }(); return variant; } diff --git a/backends/vulkan/test/custom_ops/CMakeLists.txt b/backends/vulkan/test/custom_ops/CMakeLists.txt index f29b518ec06..ad507345b23 100644 --- a/backends/vulkan/test/custom_ops/CMakeLists.txt +++ b/backends/vulkan/test/custom_ops/CMakeLists.txt @@ -107,5 +107,6 @@ if(TARGET vulkan_backend) add_operator_prototype(test_q8ta_conv2d_dw) add_operator_prototype(test_mm) add_operator_prototype(test_coopmat_probe) + add_operator_prototype(test_llama_microbench) add_operator_prototype(test_coopmat_linear_bench) endif() diff --git a/backends/vulkan/test/custom_ops/test_llama_microbench.cpp b/backends/vulkan/test/custom_ops/test_llama_microbench.cpp new file mode 100644 index 00000000000..99693ebeac8 --- /dev/null +++ b/backends/vulkan/test/custom_ops/test_llama_microbench.cpp @@ -0,0 +1,1595 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +// Unified Llama microbenchmark: the merge of the three previous harnesses +// (test_coopmat_linear_bench, test_llama_baseline_bench, +// test_sdpa_coopmat_bench) into one binary, at the real e2e dispatch shapes +// of Llama 3.1 8B / 3.2 3B / 3.2 1B (2048-token prefill, single-token +// decode; BENCHMARKING.md's ctx3072 PTEs). +// +// Suites (all run when no suite flag is given): +// --linear coopmat-vs-tiled for the two int4 linear types: +// 4w = linear_q4gsw (weight-only int4) +// 8da4w = linear_dq8ca_q4gsw (dyn-act int8 x int4 weight) +// Texture3D+Half output selects the tiled baseline; Buffer+Half +// lets QuantizedLinear.cpp's gate pick the _coopmat (WMMA) +// shader at prefill. At decode (M=1) is_gemv_case +// short-circuits to the "_coop" gemv shader for BOTH storages. +// --baseline the same linear cases run with ET_VK_FORCE_TILED_LINEAR=1 +// (specs/001's no-WMMA baseline): the buffer rows give the +// forced-tiled reference on the SAME storage the coopmat +// shader uses, isolating the algorithm from the storage type +// (specs/004). +// --sdpa llama.custom_sdpa.default tiled-vs-coopmat at each model's +// real attention shape (specs/010/021): prefill S=2048/ctx=3072 +// and decode S=1/ctx=3072/input_pos=3071 (the single most +// expensive real decode step). SDPA coopmat is default-on in +// this tree; ET_VK_DISABLE_COOPMAT is the kill switch, toggled +// per-case here to measure both variants. Decode never +// considers coopmat (is_gemv), so only tiled is measured there. +// +// Other flags: +// --model= only run models whose name contains +// --correctness-only run just the linear correctness matrix, skip perf +// --skip-correctness skip the linear correctness gate before perf +// --list print every case that would run (with sizes), no GPU +// --help +// +// Matching the real exported model: +// - per-model linear (K,N) from each checkpoint's params.json. lm_head +// (K,128256) is excluded per specs/021's explicit decision (largest and +// wildly-variable dispatch; QueryPool-race / GPU-reset trigger). +// - linear prefill M=2048, decode M=1; group_size 32 (`--group_size 32` / +// 8da4w default), coopmat-eligible for both ops' tile geometries. +// - rank-3 [1, M, K] activations, never squeezed (specs/003) -- admitted +// to the coopmat path by specs/009's leading-dims==1 relaxation. +// +// Output: one specs/021-schema "RESULT,..." line per case streamed during +// the run (shared 12 fields, then suite-specific extras -- linear/baseline: +// storage, M; sdpa: num_kv_heads, toggle), then a report: raw-results +// table, per-site WMMA speedups (coopmat vs tiled -- and vs the forced-tiled +// buffer baseline when --baseline ran), and geomeans per scheme/model/suite +// plus an overall geomean. +// +// Perf cases run one execute_test_cases() call each (specs/021 Decision 8's +// pattern) so peak host memory stays bounded by a single case's tensors +// (the 8B M=2048 FFN cases are ~115MB each). Perf cases are perf-only: +// bench_reference rejects their sizes -> SKIPPED (correctness is covered by +// the small deterministic matrix run as a gate before the sweep). + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "utils.h" + +using namespace executorch::vulkan::prototyping; +using namespace vkcompute; + +namespace { + +// ======================= shared: records + output ======================= + +struct Record { + std::string suite; // "linear" | "baseline" | "sdpa" + std::string model; + std::string scheme; // "4w" | "8da4w" | "" (sdpa) + std::string regime; // "prefill" | "decode" + // linear/baseline: op label (wq_wo, ...); sdpa: sub-shader (qk/av/total) + std::string op; + std::string storage; // linear/baseline: texture3d | buffer + // linear/baseline: dispatched kernel class (tiled/coop/coopmat/crashed); + // sdpa: which ET_VK_DISABLE_COOPMAT setting produced this row + std::string variant; + std::string kernel; // full dispatched shader name (informational) + int64_t M = 0; // sdpa: seq_len + int64_t K = 0; // sdpa: head_dim + int64_t N = 0; // sdpa: num_heads + int64_t kv = 0; // sdpa only: num_kv_heads + float mean_us = -1.0f; + float stdev_us = -1.0f; + // linear/baseline only: the linear_* shader's own per-invocation time. + // mean_us is OP-level (all unfiltered dispatches -- for 8da4w that + // includes the activation quantize_and_pack shader, real per-op e2e + // overhead); kernel_us isolates the linear kernel itself so the two + // schemes' shader-level numbers stay comparable. -1 for sdpa rows. + float kernel_us = -1.0f; + float gflops = -1.0f; // no SDPA meaning (-1 sentinel, per specs/021) + std::string dispatch = "not_applicable"; + std::string correctness = "SKIPPED"; + bool ok = false; +}; + +std::vector g_records; + +// specs/021 (research.md Decision 1): shared unified RESULT,... line -- +// 12 shared fields, then suite-specific extras. +void emit(const Record& r) { + g_records.push_back(r); + std::cout << "RESULT," << r.suite << "," << r.model << "," << r.scheme << "," + << r.regime << ","; + if (r.suite == "sdpa") { + std::cout << r.op << "," << r.K << "," << r.N << "," << r.mean_us << "," + << r.stdev_us << ",-1," << r.dispatch << ",SKIPPED," << r.kv + << "," << r.variant << "\n"; + } else { + // r.kernel (the full dispatched shader name) is appended last, after the + // specs/021 fields, so existing parsers keep working. It is the ONLY + // field that identifies WHICH tile variant ran: r.variant is + // kernel_class(), which collapses every shader to coopmat/coop/tiled, + // and r.dispatch is derived from r.variant. Without this a tile-sweep + // driver cannot tell that an unrecognized ET_VK_*_COOPMAT_VARIANT token + // silently fell back to the default kernel. + std::cout << r.variant << "," << r.K << "," << r.N << "," << r.mean_us + << "," << r.stdev_us << "," << r.gflops << "," << r.dispatch + << "," << r.correctness << "," << r.storage << "," << r.M << "," + << r.kernel_us << "," << r.kernel << "\n"; + } +} + +float geomean(const std::vector& v) { + if (v.empty()) { + return 0.0f; + } + double acc = 0.0; + for (float x : v) { + acc += std::log(static_cast(x)); + } + return static_cast(std::exp(acc / static_cast(v.size()))); +} + +// ===================== linear / baseline suites ===================== + +struct LinearConfig { + int64_t M; + int64_t K; + int64_t N; + int64_t group_size; // only meaningful for 4-bit + std::string op_name; + // 0 = rank-2 input/output ({M,K}/{M,N}, the correctness matrix below). + // >=1 = rank-3 ({batch,M,K}/{batch,M,N}) -- the real exported model's + // rank-3, batch=1 activations (specs/003, never squeezed), admitted to + // the coopmat path by specs/009's guard relaxation. All perf cases run + // this way; kRank3CorrectnessShapes carries the correctness coverage. + int64_t batch = 0; + // Perf-case labeling; empty for correctness cases. + std::string model; + std::string regime; // "prefill" | "decode" + std::string op_label; // "wq_wo", "wk_wv", ... +}; + +bool is_dq8ca(const std::string& op) { + return op.find("dq8ca") != std::string::npos; +} +bool is_4bit(const std::string& op) { + return op.find("q4gsw") != std::string::npos; +} + +// Build one test case for the given op at (storage, half dtype), no bias. +TestCase make_linear_case(const LinearConfig& cfg, utils::StorageType storage) { + const vkapi::ScalarType dt = vkapi::kHalf; + TestCase tc; + const std::string storage_str = + (storage == utils::kTexture3D) ? "Texture3D" : "Buffer"; + const std::string prefix = cfg.model.empty() + ? "" + : cfg.model + "_" + cfg.regime + "_" + cfg.op_label + "_"; + tc.set_name( + prefix + cfg.op_name + "_M" + std::to_string(cfg.M) + "_K" + + std::to_string(cfg.K) + "_N" + std::to_string(cfg.N) + + (cfg.batch > 0 ? "_rank3batch" + std::to_string(cfg.batch) : "") + "_" + + storage_str); + tc.set_operator_name("et_vk." + cfg.op_name + ".default"); + + const std::vector input_sizes = cfg.batch > 0 + ? std::vector{cfg.batch, cfg.M, cfg.K} + : std::vector{cfg.M, cfg.K}; + ValueSpec input( + input_sizes, dt, storage, utils::kWidthPacked, DataGenType::RANDINT); + + // dynamic per-row activation scale/zp (dq8ca only) + ValueSpec input_scale( + {1, cfg.M}, dt, storage, utils::kWidthPacked, DataGenType::RANDOM_SCALES); + input_scale.set_constant(true); + ValueSpec input_zp( + {1, cfg.M}, + vkapi::kChar, + storage, + utils::kWidthPacked, + DataGenType::RANDINT); + input_zp.set_constant(true); + + // weight + scales + sums depend on 4-bit vs 8-bit + const bool four = is_4bit(cfg.op_name); + ValueSpec qweight( + four ? std::vector{cfg.N, cfg.K / 2} + : std::vector{cfg.N, cfg.K}, + four ? vkapi::kByte : vkapi::kChar, + storage, + utils::kWidthPacked, + four ? DataGenType::RANDINT4 : DataGenType::RANDINT8); + qweight.set_constant(true); + if (four) { + qweight.set_int4(true); + } + + std::vector scales_size = four + ? std::vector{cfg.K / cfg.group_size, cfg.N} + : std::vector{cfg.N}; + ValueSpec weight_scales( + scales_size, + dt, + storage, + utils::kWidthPacked, + DataGenType::RANDOM_SCALES); + weight_scales.set_constant(true); + + ValueSpec weight_sums( + scales_size, + vkapi::kInt, + storage, + utils::kWidthPacked, + DataGenType::ZEROS); + weight_sums.set_constant(true); + if (four) { + compute_weight_sums_4bit_grouped( + weight_sums, qweight, cfg.K / cfg.group_size, cfg.N, cfg.group_size); + } else { + compute_weight_sums(weight_sums, qweight, cfg.N, cfg.K); + } + + ValueSpec group_size_spec(static_cast(cfg.group_size)); + + ValueSpec bias({cfg.N}, dt, storage, utils::kWidthPacked, DataGenType::ZEROS); + bias.set_constant(true); + bias.set_none(true); + + const std::vector output_sizes = cfg.batch > 0 + ? std::vector{cfg.batch, cfg.M, cfg.N} + : std::vector{cfg.M, cfg.N}; + ValueSpec output( + output_sizes, dt, storage, utils::kWidthPacked, DataGenType::ZEROS); + + // assemble per op signature + if (cfg.op_name == "linear_q4gsw") { + tc.add_input_spec(input); + tc.add_input_spec(qweight); + tc.add_input_spec(weight_scales); + tc.add_input_spec(group_size_spec); + tc.add_input_spec(bias); + } else if (cfg.op_name == "linear_dq8ca_q4gsw") { + tc.add_input_spec(input); + tc.add_input_spec(input_scale); + tc.add_input_spec(input_zp); + tc.add_input_spec(qweight); + tc.add_input_spec(weight_sums); + tc.add_input_spec(weight_scales); + tc.add_input_spec(group_size_spec); + tc.add_input_spec(bias); + } + tc.add_output_spec(output); + return tc; +} + +// ---- correctness reference for both ops; oversized shapes (the perf +// cases) throw -> framework marks them SKIPPED. For dq8ca the activation +// quant round-trip (round(x/scale)+zp) is mirrored in fp32; this is exact +// (not just close) for the correctness data below, which uses scale=1/16, +// zp=0 and activations that are multiples of 1/16, so fp16-vs-fp32 +// divergence cannot occur. ---- +std::vector as_f(const ValueSpec& s) { + if (s.dtype == vkapi::kFloat) { + return s.get_float_data(); + } + const auto& h = s.get_half_data(); + std::vector o(h.size()); + for (size_t i = 0; i < h.size(); ++i) { + o[i] = half_to_float(h[i]); + } + return o; +} +void bench_reference(TestCase& tc) { + const std::string op = tc.operator_name(); + const bool dq8ca = op.find("dq8ca") != std::string::npos; + const bool four = op.find("q4gsw") != std::string::npos; + const ValueSpec& in = tc.inputs()[0]; + ValueSpec& out = tc.outputs()[0]; + // Rank-agnostic: reads the trailing two dims, so a rank-3 [batch, M, K] + // input (batch=1, specs/009) is handled identically to plain [M, K] -- + // the reference matmul below only ever needs (M, K, N), never the batch. + const auto is = in.get_tensor_sizes(); + const int64_t M = is[is.size() - 2], K = is[is.size() - 1]; + const int64_t N = out.get_tensor_sizes().back(); + // M/N stay capped at 256 (the perf-sweep shapes reuse this same function + // and go up to M=2048/N=14336 -- an O(M*N*K) CPU reference at that size + // would take far too long and isn't the point of a perf case anyway). + // K's cap is raised for specs/014-m5-linear-coopmat-retune's FR-008 + // production-K correctness cases (K=2048/4096, M/N still <=256): without + // this, those cases silently throw here and get marked SKIPPED, giving a + // false impression of "validated" when no reference was ever computed. + if (M > 256 || N > 256 || K > 4096) { + throw std::invalid_argument("ref: too big"); + } + // input layouts: weight-only = {in, w, w_scales, [group], bias}; + // dq8ca = {in, in_scale, in_zp, w, w_sums, w_scales, [group], bias} + const ValueSpec& w = tc.inputs()[dq8ca ? 3 : 1]; + const ValueSpec& sc = tc.inputs()[dq8ca ? 5 : 2]; + const int64_t group = four ? tc.inputs()[dq8ca ? 6 : 3].get_int_value() : K; + const ValueSpec& bias = tc.inputs()[dq8ca ? (four ? 7 : 6) : (four ? 4 : 3)]; + const bool has_bias = !bias.is_none(); + + const std::vector inf = as_f(in); + const std::vector scf = as_f(sc); + const std::vector bf = has_bias ? as_f(bias) : std::vector(); + const std::vector in_scale = + dq8ca ? as_f(tc.inputs()[1]) : std::vector(); + const std::vector& in_zp = + dq8ca ? tc.inputs()[2].get_int8_data() : std::vector(); + const std::vector& w4 = + four ? w.get_uint8_data() : std::vector(); // [N, K/2] nibbles + const std::vector& w8 = + four ? std::vector() : w.get_int8_data(); // [N, K] + + auto& ref = out.get_ref_float_data(); + ref.resize(M * N); + for (int64_t m = 0; m < M; ++m) { + const float s_in = dq8ca ? in_scale[m] : 1.0f; + const int zp = dq8ca ? int(in_zp[m]) : 0; + for (int64_t n = 0; n < N; ++n) { + float acc = 0.0f; + for (int64_t k = 0; k < K; ++k) { + float a = inf[m * K + k]; + if (dq8ca) { + float q = std::round(a / s_in) + float(zp); + q = std::min(std::max(q, -128.0f), 127.0f); + a = q - float(zp); + } + int wv; + if (four) { + const uint8_t byte = w4[n * (K / 2) + k / 2]; + const int nib = (k & 1) ? ((byte >> 4) & 0xF) : (byte & 0xF); + wv = nib - 8; + } else { + wv = w8[n * K + k]; + } + const float w_scale = four ? scf[(k / group) * N + n] : scf[n]; + acc += a * float(wv) * w_scale; + } + float r = dq8ca ? acc * s_in : acc; + if (has_bias) { + r += bf[n]; + } + ref[m * N + n] = r; + } + } +} + +// Real per-model linear weight shapes (K,N), from each checkpoint's +// params.json -- the same table specs/001's shapes.json carries. wq/wo, +// wk/wv, and w1/w3 share a (K,N) within each model, so each unique dispatch +// shape is measured once and labeled with both ops. lm_head is excluded +// (see file header). +struct OpShape { + const char* op_label; + int64_t K; + int64_t N; +}; +struct LinearModel { + const char* model; + std::vector ops; +}; +const std::vector kLinearModels = { + {"llama-3.1-8b", + {{"wq_wo", 4096, 4096}, + {"wk_wv", 4096, 1024}, + {"w1_w3", 4096, 14336}, + {"w2", 14336, 4096}}}, + {"llama-3.2-3b", + {{"wq_wo", 3072, 3072}, + {"wk_wv", 3072, 1024}, + {"w1_w3", 3072, 8192}, + {"w2", 8192, 3072}}}, + {"llama-3.2-1b", + {{"wq_wo", 2048, 2048}, + {"wk_wv", 2048, 512}, + {"w1_w3", 2048, 8192}, + {"w2", 8192, 2048}}}, +}; +const std::vector> kSchemes = { + {"4w", "linear_q4gsw"}, + {"8da4w", "linear_dq8ca_q4gsw"}}; +// Real regimes: prefill dispatches every linear at M=2048 (the full prompt); +// each of the 1024 decode steps dispatches at M=1, independent of position. +const std::vector> kLinearRegimes = { + {"prefill", 2048}, + {"decode", 1}}; +// Linear-weight quantization group size. MUST match the pte being modelled: +// the export recipe uses quantization.group_size=128 (setup/README.md:247,429; +// specs/036 protocol.md "this box's buffer ptes: 128"). The previous +// hardcoded 32 was the EMBEDDING group (embedding_quantize=4,32) and is a +// different tensor; at 32 the runtime's group_size %% tile_k == 0 check +// (QuantizedLinear.cpp) rejects every tile_k>32 variant, silently falling +// back to the tiled kernel -- 71 of 160 dbuf4 tokens, i.e. the whole +// tile_k in {64,128} subspace. Overridable so a sweep can state it. +int64_t g_group = 128; +constexpr int kWarmupRuns = 3; +constexpr int kTimedRuns = 5; + +// Case selection filters (specs/041 tile sweep). A tile-variant token only +// affects ONE (scheme, storage) cell: an ET_VK_Q4GSW_COOPMAT_VARIANT token +// changes 4w+buffer, an ET_VK_DQ8CA_COOPMAT_VARIANT token changes +// 8da4w+buffer. The texture3d rows are the tiled baseline and the other +// scheme is untouched, so running them per token is pure waste -- across a +// 160-token sweep restricting to the affected cell cuts wall clock ~4x. +// Empty string = no filter (the pre-existing all-cases behaviour). +// --regime is the other half of the saving: at M=1 the linear op takes the +// is_gemv short-circuit and dispatches linear_*_coop, NOT the tsweep coopmat +// variant (verified on device: decode rows report dispatch=not_applicable and +// kernel=linear_q4gsw_coop_...). A tile token therefore cannot change a decode +// row at all, so sweeping it over decode measures the same kernel 160 times. +struct CaseFilter { + std::string model; // substring match on model name + std::string scheme; // "4w" | "8da4w" + std::string storage; // "buffer" | "texture3d" + std::string regime; // "prefill" | "decode" +}; + +bool regime_selected(const CaseFilter& f, const char* regime) { + return f.regime.empty() || f.regime == regime; +} + +bool model_selected(const CaseFilter& f, const char* model) { + return f.model.empty() || + std::string(model).find(f.model) != std::string::npos; +} + +bool scheme_selected(const CaseFilter& f, const char* scheme_label) { + return f.scheme.empty() || f.scheme == scheme_label; +} + +bool storage_selected(const CaseFilter& f, utils::StorageType st) { + if (f.storage.empty()) { + return true; + } + return f.storage == (st == utils::kTexture3D ? "texture3d" : "buffer"); +} + +// Builds one deterministic, well-conditioned correctness case (POSITIVE +// data, no fp16 cancellation -- see generate_correctness_cases) for the +// given op/shape/storage. +TestCase make_deterministic_correctness_case( + const LinearConfig& cfg, + const std::string& op, + utils::StorageType st) { + const bool dq = is_dq8ca(op); + const bool four = is_4bit(op); + TestCase t = make_linear_case(cfg, st); + auto& hin = t.inputs()[0].get_half_data(); + for (size_t i = 0; i < hin.size(); ++i) { + hin[i] = float_to_half(0.5f + 0.125f * float(i % 8)); + } + const size_t w_idx = dq ? 3 : 1; + if (four) { + auto& wq = t.inputs()[w_idx].get_uint8_data(); + const uint8_t kPos[6] = {0x99, 0xAA, 0xBB, 0xCC, 0xDD, 0xEE}; + for (size_t i = 0; i < wq.size(); ++i) { + wq[i] = kPos[i % 6]; + } + } else { + auto& wq = t.inputs()[w_idx].get_int8_data(); + for (size_t i = 0; i < wq.size(); ++i) { + wq[i] = int8_t(1 + (i % 6)); + } + } + if (dq) { + auto& hs = t.inputs()[1].get_half_data(); + std::fill(hs.begin(), hs.end(), float_to_half(0.0625f)); + auto& zp = t.inputs()[2].get_int8_data(); + std::fill(zp.begin(), zp.end(), int8_t(0)); + // weights were overwritten above -> recompute the sums + if (four) { + compute_weight_sums_4bit_grouped( + t.inputs()[4], + t.inputs()[w_idx], + cfg.K / cfg.group_size, + cfg.N, + cfg.group_size); + } else { + compute_weight_sums(t.inputs()[4], t.inputs()[w_idx], cfg.N, cfg.K); + } + } + t.set_abs_tolerance(0.5f); + t.set_rel_tolerance(0.05f); + return t; +} + +// Correctness: small aligned cases for both ops; the buffer case fires the +// coopmat shader, validated against the fp32 reference (the perf cases are +// rejected by it). POSITIVE well-conditioned data (no fp16 cancellation): +// activations are multiples of 1/16 in [0.5,1.375]; int4 nibbles in {9..14} +// (-> weight +1..+6). For dq8ca the per-row activation scale is forced to +// 1/16 with zp=0 so the dynamic int8 quant round-trip is EXACT in both fp16 +// and fp32 and the fp32 reference is valid. fp16~=fp32 throughout, so a +// tight tolerance validates shader structure (catches zero-subtile bugs) +// while ignoring benign fp16 noise. Texture3D = tiled, Buffer = coopmat. +// Shapes align to BOTH coopmat geometries (64x64x32 legacy, 128x128x16 +// double-buffered); the second shape dispatches a multi-workgroup grid for +// both, covering the gl_WorkGroupID-derived tile offsets in the store +// address math. +// Production-K cases (specs/014, FR-008): fp16-accumulation drift grows +// with the K-length of the reduction, so a shader change to the accumulator +// path can pass at small K and still diverge at production K -- the +// K=2048/4096 entries close that gap. (These correctness rows keep their +// original group sizes; the perf sweep's group_size is g_group.) +std::vector generate_correctness_cases(const CaseFilter& filter) { + std::vector cases; + static const std::vector kCorrectnessShapes = { + {64, 128, 64, 64, ""}, + {128, 256, 128, 64, ""}, + {128, 128, 128, 64, ""}, + {256, 256, 256, 64, ""}, + // Discriminators for the tiled-texture cube-shape failure: + {128, 128, 256, 64, ""}, // M == K only + {256, 128, 128, 64, ""}, // K == N only + {64, 128, 256, 64, ""}, // K > M, K < N + {256, 128, 64, 64, ""}, // K < M, K > N + // Production-K (FR-008): + {128, 2048, 128, 128, ""}, + {128, 4096, 128, 128, ""}}; + for (const auto& scheme : kSchemes) { + if (!scheme_selected(filter, scheme.first)) { + continue; + } + for (const auto& shape : kCorrectnessShapes) { + LinearConfig cfg{ + shape.M, shape.K, shape.N, shape.group_size, scheme.second}; + for (auto st : {utils::kTexture3D, utils::kBuffer}) { + if (!storage_selected(filter, st)) { + continue; + } + cases.push_back( + make_deterministic_correctness_case(cfg, scheme.second, st)); + } + } + } + // Rank-3, batch=1 correctness cases (specs/009): the real exported + // model's linear activations are rank-3 [1, M, K]. The perf sweep runs + // rank-3 too, but perf-only -- these are the cases that actually validate + // the shape against the fp32 reference, at Buffer storage (the storage + // the coopmat path requires); Texture3D+rank-3 exercises the pre-existing + // tiled path, so is not repeated here. The K=4096 entry is the + // production-K rank-3 case (FR-008). + static const std::vector kRank3CorrectnessShapes = { + {128, 128, 128, 64, "", /*batch=*/1}, + {128, 4096, 128, 128, "", /*batch=*/1}}; + // Rank-3 was buffer-only because coopmat was buffer-only. With texture IO + // (ET_VK_TEXTURE_COOPMAT) it must also run at texture3d: the rank-2 cases + // above fall back to tiled at large tiles, so rank-3 is the ONLY correctness + // case that exercises the texture coopmat epilogue at MMAS_PER_SG_M > 1 -- + // i.e. the dynamically-indexed result[i][j] drain that specs/040 flags as + // the Xclipse/PAL risk. Validating only small tiles would miss it entirely. + const utils::StorageType rank3_storage = + filter.storage == "texture3d" ? utils::kTexture3D : utils::kBuffer; + for (const auto& scheme : kSchemes) { + if (!scheme_selected(filter, scheme.first) || + !storage_selected(filter, rank3_storage)) { + continue; + } + for (const auto& shape : kRank3CorrectnessShapes) { + LinearConfig cfg{ + shape.M, + shape.K, + shape.N, + shape.group_size, + scheme.second, + shape.batch}; + cases.push_back(make_deterministic_correctness_case( + cfg, scheme.second, rank3_storage)); + } + } + return cases; +} + +int64_t flop_calc(const TestCase& tc) { + const auto& in = tc.inputs()[0].get_tensor_sizes(); + const auto& out = tc.outputs()[0].get_tensor_sizes(); + const int64_t M = in[in.size() - 2], K = in[in.size() - 1], N = out.back(); + return 2 * M * N * K; // MAC = 2 flops +} + +// The result's kernel_name is the test-case name; the dispatched shader +// names are in the per-shader timings (dq8ca cases also run a +// quantize_and_pack shader, so pick the linear_* one). +std::string linear_kernel(const BenchmarkResult& r) { + std::string name = r.get_kernel_name(); + for (const auto& st : r.get_shader_timings()) { + if (st.shader_name.find("linear_") != std::string::npos) { + name = st.shader_name; + } + } + return name; +} + +// Per-invocation time of the linear_* shader alone. ShaderTiming holds one +// iter_timings_us entry per dispatch (chained dispatches included), so its +// get_avg_time_us() is already per-invocation -- unlike the case-level +// mean_us, which sums every unfiltered dispatch (for dq8ca that adds the +// activation quantize_and_pack shader). +float linear_kernel_us(const BenchmarkResult& r) { + float us = -1.0f; + for (const auto& st : r.get_shader_timings()) { + if (st.shader_name.find("linear_") != std::string::npos) { + us = st.get_avg_time_us(); + } + } + return us; +} + +std::string kernel_class(const std::string& kernel) { + // _coopmat must be checked before _coop (substring). + if (kernel.find("_coopmat") != std::string::npos) { + return "coopmat"; + } + if (kernel.find("_coop") != std::string::npos) { + return "coop"; + } + return "tiled"; +} + +// Runs the linear correctness matrix and the specs/009 rank-3 +// dispatch+correctness verdict. Returns false on any failure. Must run +// WITHOUT ET_VK_FORCE_TILED_LINEAR set -- the buffer cases exist to fire +// and validate the coopmat shader. +// +// One execute_test_cases() call per case, each wrapped in try/catch: the +// framework throws on a numeric validation failure (utils.cpp's +// execute_test_cases, outside its own try/catch), so a batched call would +// die at the FIRST failing case and never enumerate the rest. Per-case +// execution turns that throw into one recorded failure and keeps going -- +// the whole point of the gate is to list everything that broke. (Costs the +// cross-case reference cache, but every shape here is small.) +bool run_linear_correctness(const CaseFilter& filter) { + unsetenv("ET_VK_FORCE_TILED_LINEAR"); + std::vector results; + std::vector failed_names; + for (auto& tc : generate_correctness_cases(filter)) { + try { + auto res = execute_test_cases( + [&tc]() { return std::vector{tc}; }, + flop_calc, + "LlamaMicrobenchCorrectness", + kWarmupRuns, + kTimedRuns, + bench_reference); + if (!res.empty()) { + if (res[0].get_correctness_status() == CorrectnessStatus::FAILED) { + failed_names.push_back(tc.name()); + } + results.push_back(res[0]); + } + } catch (const std::exception& e) { + failed_names.push_back(tc.name()); + std::cout << "[correctness] " << tc.name() << " FAILED: " << e.what() + << "\n"; + } + } + bool all_ok = failed_names.empty(); + // Rank-3, batch=1 dispatch + correctness verdict: numeric PASS alone + // doesn't prove the coopmat path actually ran -- the tiled fallback would + // numerically pass too. Explicitly confirm the dispatched kernel name. + for (const auto& r : results) { + if (r.get_kernel_name().find("_rank3batch") == std::string::npos) { + continue; + } + const std::string shader_name = linear_kernel(r); + const bool fired = shader_name.find("coopmat") != std::string::npos; + const bool ok = + fired && r.get_correctness_status() == CorrectnessStatus::PASSED; + all_ok = all_ok && ok; + std::cout << "[rank3 batch=1] " << r.get_kernel_name() << " -> " + << shader_name + << (fired ? " (coopmat dispatched)" + : " (NOT coopmat -- fallback)") + << ", correctness=" + << (r.get_correctness_status() == CorrectnessStatus::PASSED + ? "PASSED" + : (r.get_correctness_status() == CorrectnessStatus::FAILED + ? "FAILED" + : "SKIPPED")) + << "\n"; + } + if (!failed_names.empty()) { + std::cout << "[correctness] " << failed_names.size() + << " case(s) FAILED:\n"; + for (const auto& n : failed_names) { + std::cout << " " << n << "\n"; + } + } + if (!all_ok) { + std::cout << "[correctness] FAILED -- numeric failure(s) and/or a rank-3 " + "case did not dispatch coopmat\n"; + } + return all_ok; +} + +struct PerfCase { + LinearConfig cfg; + utils::StorageType storage; +}; +std::vector generate_linear_perf_cases(const CaseFilter& filter) { + std::vector cases; + for (const auto& scheme : kSchemes) { + if (!scheme_selected(filter, scheme.first)) { + continue; + } + for (const auto& model : kLinearModels) { + if (!model_selected(filter, model.model)) { + continue; + } + for (const auto& regime : kLinearRegimes) { + if (!regime_selected(filter, regime.first)) { + continue; + } + for (const auto& shape : model.ops) { + LinearConfig cfg{ + regime.second, + shape.K, + shape.N, + g_group, + scheme.second, + /*batch=*/1, + model.model, + regime.first, + shape.op_label}; + if (storage_selected(filter, utils::kTexture3D)) { + cases.push_back({cfg, utils::kTexture3D}); // tiled/gemv baseline + } + if (storage_selected(filter, utils::kBuffer)) { + cases.push_back({cfg, utils::kBuffer}); // coopmat (gate-permitting) + } + } + } + } + } + return cases; +} + +// Runs the linear perf sweep as either the "linear" suite (coopmat enabled) +// or the "baseline" suite (ET_VK_FORCE_TILED_LINEAR=1 for every case -- +// specs/001's no-WMMA baseline; its buffer rows are the forced-tiled +// reference on the same storage the coopmat shader uses). One +// execute_test_cases() call per case (see file header); a case-local +// failure is recorded as a crashed row and must not take down the sweep. +void run_linear_suite(const std::string& suite, const CaseFilter& filter) { + const bool force_tiled = suite == "baseline"; + if (force_tiled) { + setenv("ET_VK_FORCE_TILED_LINEAR", "1", /*overwrite=*/1); + } else { + unsetenv("ET_VK_FORCE_TILED_LINEAR"); + } + for (const auto& pc : generate_linear_perf_cases(filter)) { + const LinearConfig& cfg = pc.cfg; + Record rec; + rec.suite = suite; + rec.model = cfg.model; + rec.scheme = is_dq8ca(cfg.op_name) ? "8da4w" : "4w"; + rec.regime = cfg.regime; + rec.op = cfg.op_label; + rec.storage = pc.storage == utils::kTexture3D ? "texture3d" : "buffer"; + rec.M = cfg.M; + rec.K = cfg.K; + rec.N = cfg.N; + rec.variant = "crashed"; + rec.kernel = "CRASHED"; + rec.correctness = "SKIPPED"; // perf-only; see bench_reference + try { + TestCase tc = make_linear_case(cfg, pc.storage); + auto res = execute_test_cases( + [&tc]() { return std::vector{tc}; }, + flop_calc, + "LlamaMicrobench", + kWarmupRuns, + kTimedRuns, + bench_reference); + if (!res.empty()) { + rec.mean_us = res[0].get_avg_time_us(); + rec.stdev_us = res[0].get_std_dev_us(); + rec.kernel = linear_kernel(res[0]); + rec.kernel_us = linear_kernel_us(res[0]); + rec.variant = kernel_class(rec.kernel); + rec.gflops = rec.mean_us > 0 + ? (2.0f * cfg.M * cfg.N * cfg.K) / (rec.mean_us * 1e3f) + : -1.0f; + rec.ok = true; + } + } catch (const std::exception& e) { + std::cerr << "WARNING: case '" << cfg.model << " " << cfg.regime << " " + << cfg.op_label << " " << cfg.op_name << " " << rec.storage + << "' threw: " << e.what() + << " -- recorded as CRASHED, continuing\n"; + } + // Dispatch expectation: the linear suite's prefill buffer rows must be + // coopmat ("confirmed"/"fallback_tiled" -- specs/021 FR-006 semantics); + // everywhere else no coopmat is possible (texture, decode's is_gemv + // short-circuit, or the baseline suite's forced-tiled), so the status + // is not_applicable -- except a coopmat kernel showing up there, which + // is a real anomaly. + if (rec.ok) { + // texture3d+coopmat is expected too once ET_VK_TEXTURE_COOPMAT is set; + // without this the texture runs report unexpected_coopmat and the binary + // exits nonzero despite dispatching exactly what was asked for. + static const bool tex_coopmat = + std::getenv("ET_VK_TEXTURE_COOPMAT") != nullptr; + const bool expects_coopmat = suite == "linear" && + rec.regime == "prefill" && + (rec.storage == "buffer" || + (tex_coopmat && rec.storage == "texture3d")); + if (expects_coopmat) { + rec.dispatch = + rec.variant == "coopmat" ? "confirmed" : "fallback_tiled"; + } else { + rec.dispatch = + rec.variant == "coopmat" ? "unexpected_coopmat" : "not_applicable"; + } + } else { + rec.dispatch = "crashed"; + } + emit(rec); + } + if (force_tiled) { + unsetenv("ET_VK_FORCE_TILED_LINEAR"); + } +} + +// ============================ sdpa suite ============================ + +struct SdpaModel { + const char* name; + int64_t head_dim; + int64_t num_heads; + int64_t num_kv_heads; +}; +// Real per-model shapes, derived directly from each checkpoint's params.json +// (dim / n_heads), matching specs/010 research.md Decision 5. +const std::vector kSdpaModels = { + {"llama-3.1-8b", 128, 32, 8}, + {"llama-3.2-3b", 128, 24, 8}, + {"llama-3.2-1b", 64, 32, 8}, +}; + +// specs/021: real e2e regimes. context_len=3072 for BOTH regimes: the real +// ctx3072 PTEs (2048 prefill + 1024 decode) allocate the KV cache at +// max_context_length up front, so even the prefill step's cache tensor -- +// and therefore the attention shaders' strides/access pattern -- is sized +// 3072, not 2048. input_pos=3071 is the single most expensive real decode +// step (attends over the fullest cache). SDPA.cpp's is_gemv gate means the +// coopmat toggle has no effect at decode. +struct SdpaRegime { + const char* regime; + int64_t seq_len; // this step's query / newly-written KV length + int64_t context_len; // KV cache buffer size + int64_t input_pos; // symint value +}; +const std::vector kSdpaRegimes = { + {"prefill", 2048, 3072, 0}, + {"decode", 1, 3072, 3071}, +}; + +struct SdpaRunResult { + float mean_us; // total (qk + av) + float stdev_us; + float qk_mean_us; + float qk_stdev_us; + float av_mean_us; + float av_stdev_us; + std::vector dispatched_kernels; // from the last timed run +}; + +float mean_of(const std::vector& v) { + return std::accumulate(v.begin(), v.end(), 0.0f) / + static_cast(v.size()); +} +float stdev_of(const std::vector& v, float mean) { + if (v.size() < 2) { + return 0.0f; + } + float acc = 0.0f; + for (float x : v) { + acc += (x - mean) * (x - mean); + } + return std::sqrt(acc / static_cast(v.size() - 1)); +} + +// Fills a tensor's staging buffer with random half-precision data in +// [-1, 1]. maybe_cast_and_copy_into_staging does not support a Float->Half +// conversion (throws), so the half encoding is done host-side here and +// passed with a matching src dtype. +void fill_random(ComputeGraph& graph, const ValueRef staging, int64_t numel) { + std::vector data(numel); + for (int64_t i = 0; i < numel; ++i) { + const float v = (static_cast(std::rand()) / RAND_MAX) * 2.0f - 1.0f; + data[i] = float_to_half(v); + } + graph.maybe_cast_and_copy_into_staging( + staging, data.data(), static_cast(numel), vkapi::kHalf); +} + +// Runs one (model, coopmat-toggle, regime) case: builds +// llama.custom_sdpa.default directly via ComputeGraph (specs/010 research.md +// Decision 8: the TestCase framework has no SymInt support and this op +// family requires one). Returns the SDPA-compute-only GPU time split into +// qk (sdpa_compute_attn_weights_*) and av (sdpa_compute_out_*), plus their +// combined total -- excluding the kv-cache-update and softmax dispatches in +// between (unaccelerated, identical regardless of the coopmat toggle). +// +// SDPA coopmat is default-on in this tree; ET_VK_DISABLE_COOPMAT is the +// kill switch (read at shader-pick / graph-build time), so "tiled" here +// means running with it set. The KV cache buffer is always sized to +// `regime.context_len`, filled with random data BEFORE update_cache writes +// this step's new K/V at input_pos -- the shapes and access pattern are +// real even though the "history" isn't a genuine step-by-step prefill walk +// (only timing is measured here, not output correctness). +SdpaRunResult sdpa_run_case( + const SdpaModel& m, + bool enable_coopmat, + const SdpaRegime& regime) { + if (enable_coopmat) { + unsetenv("ET_VK_DISABLE_COOPMAT"); + } else { + setenv("ET_VK_DISABLE_COOPMAT", "1", /*overwrite=*/1); + } + + GraphConfig config; + config.enable_querypool = true; + api::context()->initialize_querypool(); + ComputeGraph graph(config); + + const int64_t batch_size = 1; + const std::vector q_sizes = { + batch_size, regime.seq_len, m.num_heads, m.head_dim}; + const std::vector new_kv_sizes = { + batch_size, regime.seq_len, m.num_kv_heads, m.head_dim}; + const std::vector cache_sizes = { + batch_size, regime.context_len, m.num_kv_heads, m.head_dim}; + + IOValueRef r_q = + graph.add_input_tensor(q_sizes, vkapi::kHalf, utils::kBuffer); + IOValueRef r_k = + graph.add_input_tensor(new_kv_sizes, vkapi::kHalf, utils::kBuffer); + IOValueRef r_v = + graph.add_input_tensor(new_kv_sizes, vkapi::kHalf, utils::kBuffer); + + const ValueRef r_input_pos_symint = graph.add_symint(regime.input_pos); + const ValueRef r_out = + graph.add_tensor(q_sizes, vkapi::kHalf, utils::kBuffer); + + const ValueRef r_k_cache = + graph.add_tensor(cache_sizes, vkapi::kHalf, utils::kBuffer); + const ValueRef r_v_cache = + graph.add_tensor(cache_sizes, vkapi::kHalf, utils::kBuffer); + const ValueRef r_dummy_out = + graph.add_tensor({1}, vkapi::kHalf, utils::kBuffer); + + VK_GET_OP_FN("update_cache.default") + (graph, {r_k.value, r_k_cache, r_input_pos_symint, r_dummy_out}); + VK_GET_OP_FN("update_cache.default") + (graph, {r_v.value, r_v_cache, r_input_pos_symint, r_dummy_out}); + VK_GET_OP_FN("llama.custom_sdpa.default") + (graph, + { + r_q.value, + r_k_cache, + r_v_cache, + r_input_pos_symint, + kDummyValueRef, // attn_mask + kDummyValueRef, // dropout_p + kDummyValueRef, // is_causal + kDummyValueRef, // scale + r_out, + }); + + graph.set_output_tensor(r_out); + graph.prepare(); + graph.prepack(); + + fill_random( + graph, + r_q.staging, + batch_size * regime.seq_len * m.num_heads * m.head_dim); + fill_random( + graph, + r_k.staging, + batch_size * regime.seq_len * m.num_kv_heads * m.head_dim); + fill_random( + graph, + r_v.staging, + batch_size * regime.seq_len * m.num_kv_heads * m.head_dim); + // Note: r_k_cache/r_v_cache are plain add_tensor() outputs (not + // IOValueRef), so there is no staging buffer to pre-fill positions + // 0..input_pos-1 with. That's fine -- only timing is measured here, and + // the attention shaders' dispatch size/access pattern depends solely on + // the cache's shape (context_len), not its contents. + + for (int i = 0; i < kWarmupRuns; ++i) { + graph.execute(); + } + + std::vector total_timings_us; + std::vector qk_timings_us; + std::vector av_timings_us; + std::vector last_dispatched; + for (int i = 0; i < kTimedRuns; ++i) { + graph.execute(); + graph.context()->querypool().extract_results(); + const auto shader_results = + graph.context()->querypool().get_shader_timestamp_data(); + + float qk_time_us = 0.0f; + float av_time_us = 0.0f; + last_dispatched.clear(); + for (const auto& r : shader_results) { + last_dispatched.push_back(r.kernel_name); + const uint64_t duration_ns = r.end_time_ns - r.start_time_ns; + if (r.kernel_name.find("sdpa_compute_attn_weights") != + std::string::npos) { + qk_time_us += static_cast(duration_ns) / 1000.0f; + } else if (r.kernel_name.find("sdpa_compute_out") != std::string::npos) { + av_time_us += static_cast(duration_ns) / 1000.0f; + } + } + qk_timings_us.push_back(qk_time_us); + av_timings_us.push_back(av_time_us); + total_timings_us.push_back(qk_time_us + av_time_us); + } + + SdpaRunResult result; + result.mean_us = mean_of(total_timings_us); + result.stdev_us = stdev_of(total_timings_us, result.mean_us); + result.qk_mean_us = mean_of(qk_timings_us); + result.qk_stdev_us = stdev_of(qk_timings_us, result.qk_mean_us); + result.av_mean_us = mean_of(av_timings_us); + result.av_stdev_us = stdev_of(av_timings_us, result.av_mean_us); + result.dispatched_kernels = last_dispatched; + return result; +} + +bool has_kernel_containing( + const std::vector& kernels, + const std::string& needle) { + for (const auto& k : kernels) { + if (k.find(needle) != std::string::npos) { + return true; + } + } + return false; +} + +void emit_sdpa_records( + const SdpaModel& m, + const SdpaRegime& regime, + const std::string& toggle, + const SdpaRunResult& r, + const std::string& dispatch) { + const struct { + const char* sub; + float mean; + float stdev; + } subs[] = { + {"qk", r.qk_mean_us, r.qk_stdev_us}, + {"av", r.av_mean_us, r.av_stdev_us}, + {"total", r.mean_us, r.stdev_us}, + }; + for (const auto& s : subs) { + Record rec; + rec.suite = "sdpa"; + rec.model = m.name; + rec.regime = regime.regime; + rec.op = s.sub; + rec.variant = toggle; + rec.M = regime.seq_len; + rec.K = m.head_dim; + rec.N = m.num_heads; + rec.kv = m.num_kv_heads; + rec.mean_us = s.mean; + rec.stdev_us = s.stdev; + rec.dispatch = dispatch; + rec.ok = true; + emit(rec); + } +} + +// Returns false if any prefill case failed to confirm coopmat dispatch. +bool run_sdpa_suite(const std::string& model_filter) { + bool all_confirmed = true; + for (const auto& m : kSdpaModels) { + if (std::string(m.name).find(model_filter) == std::string::npos) { + continue; + } + for (const auto& regime : kSdpaRegimes) { + const bool is_decode = std::string(regime.regime) == "decode"; + SdpaRunResult tiled = sdpa_run_case(m, /*enable_coopmat=*/false, regime); + // Decode: SDPA.cpp's is_gemv gate never considers coopmat -- a second + // invocation would just remeasure the identical dispatch. Skip it. + SdpaRunResult coopmat = + is_decode ? tiled : sdpa_run_case(m, /*enable_coopmat=*/true, regime); + + std::string dispatch; + if (is_decode) { + dispatch = "not_applicable"; + } else { + const bool tiled_is_tiled = + !has_kernel_containing(tiled.dispatched_kernels, "_coopmat"); + const bool qk_coopmat = has_kernel_containing( + coopmat.dispatched_kernels, "sdpa_compute_attn_weights_coopmat"); + const bool av_coopmat = has_kernel_containing( + coopmat.dispatched_kernels, "sdpa_compute_out_coopmat"); + dispatch = (tiled_is_tiled && qk_coopmat && av_coopmat) + ? "confirmed" + : "fallback_tiled"; + all_confirmed = all_confirmed && dispatch == "confirmed"; + } + + emit_sdpa_records(m, regime, "tiled", tiled, dispatch); + if (!is_decode) { + emit_sdpa_records(m, regime, "coopmat", coopmat, dispatch); + } + } + } + unsetenv("ET_VK_DISABLE_COOPMAT"); // restore the tree's default-on state + return all_confirmed; +} + +// ============================== report ============================== + +const Record* find_record( + const std::string& suite, + const std::string& model, + const std::string& scheme, + const std::string& regime, + const std::string& op, + const std::string& storage, + const std::string& variant = "") { + for (const auto& r : g_records) { + if (r.suite == suite && r.model == model && r.scheme == scheme && + r.regime == regime && r.op == op && r.storage == storage && + (variant.empty() || r.variant == variant) && r.ok) { + return &r; + } + } + return nullptr; +} + +std::string fmt_us(float us) { + if (us < 0) { + return "-"; + } + std::ostringstream ss; + ss << std::fixed << std::setprecision(1) << us; + return ss.str(); +} +std::string fmt_x(float x) { + if (x <= 0) { + return "-"; + } + std::ostringstream ss; + ss << std::fixed << std::setprecision(2) << x << "x"; + return ss.str(); +} + +// Prints the raw-results table, the per-site WMMA speedups, and the +// geomeans. Returns false if any expected coopmat site failed to speed up +// AND failed to dispatch -- dispatch anomalies, not slowness, fail the run. +void print_report(bool baseline_ran) { + print_separator(); + std::cout << "==================== RAW RESULTS ====================\n"; + std::cout << std::left << std::setw(10) << "suite" << std::setw(14) << "model" + << std::setw(7) << "scheme" << std::setw(9) << "regime" + << std::setw(7) << "op" << std::setw(11) << "storage" + << std::setw(9) << "variant" << std::setw(21) << "(M,K,N)" + << std::right << std::setw(12) << "mean_us" << std::setw(10) + << "stdev" << std::setw(10) << "kern_us" << std::setw(10) + << "GFLOP/s" << " dispatch\n"; + for (const auto& r : g_records) { + std::ostringstream shape; + shape << "(" << r.M << "," << r.K << "," << r.N << ")"; + std::cout << std::left << std::setw(10) << r.suite << std::setw(14) + << r.model << std::setw(7) << (r.scheme.empty() ? "-" : r.scheme) + << std::setw(9) << r.regime << std::setw(7) << r.op + << std::setw(11) << (r.storage.empty() ? "-" : r.storage) + << std::setw(9) << r.variant << std::setw(21) << shape.str() + << std::right << std::setw(12) << fmt_us(r.mean_us) + << std::setw(10) << fmt_us(r.stdev_us) << std::setw(10) + << fmt_us(r.kernel_us) << std::setw(10) + << (r.gflops >= 0 ? fmt_us(r.gflops) : "-") << " " << r.dispatch + << "\n"; + } + + // ---- linear WMMA speedups (prefill only; decode has no coopmat) ---- + std::vector all_wmma_speedups; + bool have_linear = false; + for (const auto& r : g_records) { + have_linear = have_linear || r.suite == "linear"; + } + if (have_linear) { + // op_x = op-level speedup (all of the op's dispatches -- for 8da4w that + // includes the activation quantize_and_pack shader, which the real + // model pays on every linear, so this is the per-op e2e gain). kern_x = + // the linear shader alone, the number that judges the WMMA kernel + // itself. For 4w the two coincide (no quantize dispatch). Geomeans use + // op_x -- the e2e-relevant quantity -- with kern_x geomeans printed + // alongside per scheme. + std::cout << "\n========== LINEAR: coopmat (WMMA) vs tiled, prefill " + "M=2048 ==========\n"; + std::cout << std::left << std::setw(7) << "scheme" << std::setw(14) + << "model" << std::setw(7) << "op" << std::setw(15) << "(K,N)" + << std::right << std::setw(12) << "tiled_tex" << std::setw(12) + << "coopmat" << std::setw(9) << "op_x" << std::setw(9) + << "kern_x"; + if (baseline_ran) { + std::cout << std::setw(14) << "tiled_buf" << std::setw(9) << "vs_buf"; + } + std::cout << " (us; op_x = whole op incl. 8da4w act-quant, kern_x = " + "linear shader only)\n"; + std::vector>> scheme_geo; + std::vector>> scheme_kern_geo; + for (const auto& scheme : kSchemes) { + std::vector scheme_speedups; + std::vector scheme_kern_speedups; + for (const auto& model : kLinearModels) { + std::vector model_speedups; + for (const auto& shape : model.ops) { + const Record* tex = find_record( + "linear", + model.model, + scheme.first, + "prefill", + shape.op_label, + "texture3d"); + // Unfiltered buffer row for display (shows the actually-dispatched + // kernel even on a fallback); coopmat-filtered row for speedups + // and geomeans, so a fallback can never contribute a bogus ratio. + const Record* buf = find_record( + "linear", + model.model, + scheme.first, + "prefill", + shape.op_label, + "buffer"); + const Record* cm = (buf && buf->variant == "coopmat") ? buf : nullptr; + const Record* base_buf = baseline_ran ? find_record( + "baseline", + model.model, + scheme.first, + "prefill", + shape.op_label, + "buffer") + : nullptr; + if (tex == nullptr && buf == nullptr) { + continue; // model filtered out + } + const float speedup = (tex && cm && cm->mean_us > 0) + ? tex->mean_us / cm->mean_us + : 0.0f; + const float kern_speedup = + (tex && cm && tex->kernel_us > 0 && cm->kernel_us > 0) + ? tex->kernel_us / cm->kernel_us + : 0.0f; + const float vs_buf = (base_buf && cm && cm->mean_us > 0) + ? base_buf->mean_us / cm->mean_us + : 0.0f; + if (speedup > 0) { + model_speedups.push_back(speedup); + all_wmma_speedups.push_back(speedup); + } + if (kern_speedup > 0) { + scheme_kern_speedups.push_back(kern_speedup); + } + std::cout << std::left << std::setw(7) << scheme.first + << std::setw(14) << model.model << std::setw(7) + << shape.op_label << std::setw(15) + << ("(" + std::to_string(shape.K) + "," + + std::to_string(shape.N) + ")") + << std::right << std::setw(12) + << fmt_us(tex ? tex->mean_us : -1.0f) << std::setw(12) + << fmt_us(buf ? buf->mean_us : -1.0f) << std::setw(9) + << fmt_x(speedup) << std::setw(9) << fmt_x(kern_speedup); + if (baseline_ran) { + std::cout << std::setw(14) + << fmt_us(base_buf ? base_buf->mean_us : -1.0f) + << std::setw(9) << fmt_x(vs_buf); + } + if (buf && buf->variant != "coopmat") { + std::cout << " ! " << buf->kernel; + } + std::cout << "\n"; + } + if (!model_speedups.empty()) { + std::cout << std::left << std::setw(7) << scheme.first + << std::setw(14) << model.model << std::setw(7) << "geo" + << std::setw(15) << "" << std::right << std::setw(12) << "" + << std::setw(12) << "" << std::setw(9) + << fmt_x(geomean(model_speedups)) << "\n"; + scheme_speedups.insert( + scheme_speedups.end(), + model_speedups.begin(), + model_speedups.end()); + } + } + scheme_geo.emplace_back(scheme.first, scheme_speedups); + scheme_kern_geo.emplace_back(scheme.first, scheme_kern_speedups); + } + for (size_t i = 0; i < scheme_geo.size(); ++i) { + if (!scheme_geo[i].second.empty()) { + std::cout << "linear " << scheme_geo[i].first + << " geomean (all models): op " + << fmt_x(geomean(scheme_geo[i].second)) << ", kernel " + << fmt_x(geomean(scheme_kern_geo[i].second)) << "\n"; + } + } + std::cout << "(! = buffer case did NOT dispatch a coopmat shader; shown " + "for reference, excluded from speedups/geomeans)\n"; + } + + // ---- sdpa WMMA speedups (prefill only) ---- + bool have_sdpa = false; + for (const auto& r : g_records) { + have_sdpa = have_sdpa || r.suite == "sdpa"; + } + if (have_sdpa) { + std::cout << "\n========== SDPA: coopmat (WMMA) vs tiled, prefill S=2048 " + "==========\n"; + std::cout << std::left << std::setw(14) << "model" << std::setw(7) << "sub" + << std::right << std::setw(12) << "tiled_us" << std::setw(12) + << "coopmat_us" << std::setw(9) << "speedup" << " dispatch\n"; + std::vector sdpa_totals; + for (const auto& m : kSdpaModels) { + for (const char* sub : {"qk", "av", "total"}) { + const Record* t = + find_record("sdpa", m.name, "", "prefill", sub, "", "tiled"); + const Record* c = + find_record("sdpa", m.name, "", "prefill", sub, "", "coopmat"); + if (t == nullptr || c == nullptr) { + continue; // model filtered out + } + const float speedup = c->mean_us > 0 ? t->mean_us / c->mean_us : 0.0f; + // Only "total" (qk+av combined) feeds the geomeans -- counting qk + // and av separately alongside it would double-weight each model. + // Only confirmed-dispatch rows count. + if (std::string(sub) == "total" && speedup > 0 && + c->dispatch == "confirmed") { + sdpa_totals.push_back(speedup); + all_wmma_speedups.push_back(speedup); + } + std::cout << std::left << std::setw(14) << m.name << std::setw(7) << sub + << std::right << std::setw(12) << fmt_us(t->mean_us) + << std::setw(12) << fmt_us(c->mean_us) << std::setw(9) + << fmt_x(speedup) << " " << c->dispatch << "\n"; + } + } + if (!sdpa_totals.empty()) { + std::cout << "sdpa geomean (total, all models): " + << fmt_x(geomean(sdpa_totals)) << "\n"; + } + } + + if (!all_wmma_speedups.empty()) { + std::cout << "\nOVERALL WMMA geomean (" << all_wmma_speedups.size() + << " sites: linear prefill shapes + sdpa prefill totals): " + << fmt_x(geomean(all_wmma_speedups)) << "\n"; + } +} + +void print_usage() { + std::cout + << "test_llama_microbench: unified Llama linear/SDPA microbenchmark\n" + " --linear run the coopmat-vs-tiled linear suite\n" + " --baseline run the forced-tiled (no-WMMA) linear suite\n" + " --sdpa run the SDPA coopmat-vs-tiled suite\n" + " (no suite flag = all three)\n" + " --model= only models whose name contains \n" + " --scheme=<4w|8da4w> only this quantization scheme\n" + " --storage= only this storage type\n" + " --regime= only this regime\n" + " --group-size= linear quant group size (default 128,\n" + " must match the pte; tile_k must divide it)\n" + " (--scheme/--storage also narrow the\n" + " correctness gate; a tile-variant token\n" + " only affects one scheme+buffer cell)\n" + " --correctness-only run just the linear correctness matrix\n" + " --skip-correctness skip the correctness gate before perf\n" + " --list print every case with its sizes, no GPU\n" + " --help this message\n"; +} + +void list_cases( + bool linear, + bool baseline, + bool sdpa, + const CaseFilter& filter) { + const std::string& model_filter = filter.model; + int n = 0; + for (const char* suite : {"linear", "baseline"}) { + if ((std::string(suite) == "linear" && !linear) || + (std::string(suite) == "baseline" && !baseline)) { + continue; + } + for (const auto& pc : generate_linear_perf_cases(filter)) { + std::cout << suite << "," << pc.cfg.model << "," + << (is_dq8ca(pc.cfg.op_name) ? "8da4w" : "4w") << "," + << pc.cfg.regime << "," << pc.cfg.op_label << "," + << (pc.storage == utils::kTexture3D ? "texture3d" : "buffer") + << ",[1," << pc.cfg.M << "," << pc.cfg.K << "]x[" << pc.cfg.K + << "," << pc.cfg.N << "],group" << pc.cfg.group_size << "\n"; + ++n; + } + } + if (sdpa) { + for (const auto& m : kSdpaModels) { + if (std::string(m.name).find(model_filter) == std::string::npos) { + continue; + } + for (const auto& regime : kSdpaRegimes) { + const bool is_decode = std::string(regime.regime) == "decode"; + for (const char* toggle : {"tiled", "coopmat"}) { + if (is_decode && std::string(toggle) == "coopmat") { + continue; + } + std::cout << "sdpa," << m.name << ",," << regime.regime << ",qk+av," + << toggle << ",S" << regime.seq_len << "_ctx" + << regime.context_len << "_pos" << regime.input_pos + << ",head" << m.head_dim << "_h" << m.num_heads << "_kv" + << m.num_kv_heads << "\n"; + ++n; + } + } + } + } + std::cout << n << " cases\n"; +} + +} // namespace + +int main(int argc, char** argv) { + bool linear = false, baseline = false, sdpa = false; + bool correctness_only = false, skip_correctness = false, list_only = false; + CaseFilter filter; + for (int i = 1; i < argc; ++i) { + const std::string arg = argv[i]; + if (arg == "--linear") { + linear = true; + } else if (arg == "--baseline") { + baseline = true; + } else if (arg == "--sdpa") { + sdpa = true; + } else if (arg == "--correctness-only") { + correctness_only = true; + } else if (arg == "--skip-correctness") { + skip_correctness = true; + } else if (arg == "--list") { + list_only = true; + } else if (arg.rfind("--model=", 0) == 0) { + filter.model = arg.substr(8); + } else if (arg.rfind("--scheme=", 0) == 0) { + filter.scheme = arg.substr(9); + if (filter.scheme != "4w" && filter.scheme != "8da4w") { + std::cerr << "--scheme must be 4w or 8da4w, got: " << filter.scheme + << "\n"; + return 2; + } + } else if (arg.rfind("--storage=", 0) == 0) { + filter.storage = arg.substr(10); + if (filter.storage != "buffer" && filter.storage != "texture3d") { + std::cerr << "--storage must be buffer or texture3d, got: " + << filter.storage << "\n"; + return 2; + } + } else if (arg.rfind("--group-size=", 0) == 0) { + g_group = std::stoll(arg.substr(13)); + } else if (arg.rfind("--regime=", 0) == 0) { + filter.regime = arg.substr(9); + if (filter.regime != "prefill" && filter.regime != "decode") { + std::cerr << "--regime must be prefill or decode, got: " + << filter.regime << "\n"; + return 2; + } + } else if (arg == "--help" || arg == "-h") { + print_usage(); + return 0; + } else { + std::cerr << "unknown flag: " << arg << "\n"; + print_usage(); + return 2; + } + } + if (!linear && !baseline && !sdpa) { + linear = baseline = sdpa = true; // default: everything + } + + if (list_only) { + list_cases(linear, baseline, sdpa, filter); + return 0; + } + + set_debugging(false); + set_print_output(false); + set_print_latencies(false); + set_use_gpu_timestamps(true); + + print_performance_header(); + std::cout << "Llama microbench (3.1 8B / 3.2 3B / 3.2 1B real e2e shapes; " + "prefill 2048 / decode 1 @ ctx3072; linear group_size=" + << g_group << "; " << kWarmupRuns << " warmup + " << kTimedRuns + << " timed runs per case)\n"; + // Device provenance: without this, thermal/DVFS drift between runs (or + // between the linear and baseline suites within one run) cannot even be + // diagnosed post hoc from a saved log. + { + const auto* adapter = api::context()->adapter_ptr(); + std::cout << "DEVICE," << adapter->device_name() + << ",timestamp_period_ns=" << adapter->timestamp_period() + << ",subgroup_size=" << adapter->subgroup_size() << ",coopmat=" + << (adapter->supports_cooperative_matrix() ? "yes" : "no") + << "\n"; + } + print_separator(); + + std::srand(0); + bool ok = true; + + // Correctness gate: validates the tiled and coopmat linear kernels + // (including the rank-3 dispatch check) before any perf time is spent. + if (correctness_only) { + return run_linear_correctness(filter) ? 0 : 1; + } + if ((linear || baseline) && !skip_correctness) { + if (!run_linear_correctness(filter)) { + std::cout << "correctness gate FAILED -- not running the perf sweep\n"; + return 1; + } + } + + if (linear) { + run_linear_suite("linear", filter); + } + if (baseline) { + run_linear_suite("baseline", filter); + } + bool sdpa_confirmed = true; + if (sdpa) { + sdpa_confirmed = run_sdpa_suite(filter.model); + } + + print_report(baseline); + + // Exit code reflects dispatch sanity, not speed: every linear-suite + // prefill buffer row must have dispatched coopmat, no coopmat may appear + // where it can't (decode/forced-tiled/texture), nothing crashed, and + // every sdpa prefill case must have confirmed coopmat dispatch. + for (const auto& r : g_records) { + if (r.dispatch == "fallback_tiled" || r.dispatch == "unexpected_coopmat" || + r.dispatch == "crashed") { + ok = false; + } + } + ok = ok && sdpa_confirmed; + if (!ok) { + std::cout << "\nOne or more cases crashed or did not dispatch the " + "expected kernel (see dispatch column) -- do not trust " + "their speedup numbers.\n"; + } + return ok ? 0 : 1; +} From bcba76da3c911288a3f3043c045e0bf09d5642b8 Mon Sep 17 00:00:00 2001 From: Yanwen Xu Date: Tue, 18 Aug 2026 13:50:44 -0700 Subject: [PATCH 07/28] Document flaky coopmat correctness failure, keep 8da4w default unchanged tsweep_dbuf4_t64x64k32g24s32 passed a correctness-first tile sweep and looked like a rep-confirmed win on 1B/3B/8B across two boards, but repeating test_llama_microbench --correctness-only back-to-back on the same board/driver with no code change showed intermittent failures (1/10 runs, most shapes wrong, no crash/error) vs. 0/6 for the current default over the same window. A single correctness-bench pass is not sufficient evidence for a coopmat tile. Net effect: no default change, comment updated to record why and to flag that the sweep's other 11 "passing" candidates are equally unverified. --- .../graph/ops/impl/QuantizedLinear.cpp | 53 ++++++++++++------- 1 file changed, 35 insertions(+), 18 deletions(-) diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp index 68b9f89ea9b..6e059ae8f5c 100644 --- a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp @@ -136,24 +136,41 @@ static const std::string& dq8ca_coopmat_variant() { // -- same geometry as the shipped buffer-storage default, resolved through // the tsweep_dbuf4 texture3d-capable shader. // - // REVERTED 2026-08-18: the M51 tile sweep's apparent winner, - // tsweep_dbuf4_t128x16k64g12s64 (11.8-17.3% faster in e2e prefill tok/s), - // was shipped as this default for a short window and then found to be - // NUMERICALLY WRONG via test_llama_microbench --correctness-only -- - // every single texture3d correctness case failed, including the rank-3 - // (real model shape) ones (~70% of elements mismatched in a structured - // per-row pattern: the first M-rows correct, everything past a boundary - // wrong). Root cause not yet isolated -- checked and ruled out - // int_input_sums undersizing (github.com/pytorch/executorch/issues/21423): - // that buffer is marked unused in both this shader and the tsweep_dbuf4 - // one, identically, so it can't be the discriminator. Most likely a - // genuine indexing bug specific to that tile's spec-resolved shape - // (M=128, N=16, K=64, 1x2 subgroup grid, sub=64) in the dbuf4 template. - // This tile (t64x32k32g12s64) IS correctness-bench-confirmed clean: 0 - // failures across all texture3d cases including rank-3, real coopmat - // dispatch confirmed via the harness's own dispatch log. Do not re-ship - // t128x16k64g12s64 (or any other untested sweep candidate) without first - // passing test_llama_microbench --scheme=8da4w --storage=texture3d + // REJECTED 2026-08-18: tsweep_dbuf4_t64x64k32g24s32 looked like a clean + // win (correctness-first sweep pick, rep-confirmed faster on 1B/3B/8B + // across two boards) and was staged as the new default, but a single-shot + // test_llama_microbench --correctness-only pass is NOT sufficient evidence + // -- repeating the identical command back-to-back on the same board/driver + // with no code change showed intermittent (flaky) correctness corruption: + // 1 catastrophic failure (most shapes wrong, silent fallback to the tiled + // path which then hit release-branch's known tiled-path bug) in 10 runs, + // vs. 0/6 for this default over the same window. No crash, no error -- + // just wrong numbers some fraction of the time. This is the same failure + // class as the q4gsw-coopmat-e2e-garbage incident (see memory): a coopmat + // dispatch can silently miscompile without any signal. A single + // correctness-bench pass during a sweep is not proof of correctness for + // that reason -- any future candidate must be run repeatedly (ideally + // 10+ back-to-back correctness-only passes with zero failures) before + // being trusted, not just once. This also means the other 11 "passing" + // candidates from the 2026-08-18 correctness-first sweep are NOT + // trustworthy either -- they were only checked once each. + // + // REVERTED 2026-08-18 (earlier incident, same day): the M51 tile sweep's + // apparent winner at the time, tsweep_dbuf4_t128x16k64g12s64 (11.8-17.3% + // faster in e2e prefill tok/s), was shipped as the default for a short + // window and then found to be NUMERICALLY WRONG via test_llama_microbench + // --correctness-only -- every single texture3d correctness case failed, + // including the rank-3 (real model shape) ones (~70% of elements + // mismatched in a structured per-row pattern: the first M-rows correct, + // everything past a boundary wrong). Root cause not isolated -- checked + // and ruled out int_input_sums undersizing + // (github.com/pytorch/executorch/issues/21423): that buffer is marked + // unused in both this shader and the tsweep_dbuf4 one, identically, so it + // can't be the discriminator. Most likely a genuine indexing bug specific + // to that tile's spec-resolved shape (M=128, N=16, K=64, 1x2 subgroup grid, + // sub=64) in the dbuf4 template. Do not re-ship t128x16k64g12s64 (or any + // other untested sweep candidate) without first passing + // test_llama_microbench --scheme=8da4w --storage=texture3d // --correctness-only clean. static const std::string variant = [] { const char* env = std::getenv("ET_VK_DQ8CA_COOPMAT_VARIANT"); From 8be6567b3adaf9f454e595def81e47aa837f8852 Mon Sep 17 00:00:00 2001 From: Yanwen Xu Date: Tue, 18 Aug 2026 14:22:29 -0700 Subject: [PATCH 08/28] [ET-VK] Port SDPA cooperative-matrix (WMMA) shaders to release/1.4 Brings over the SDPA QK^T/AV coopmat shaders and SDPA.cpp gating from dev (originally bundled with the linear coopmat port in 8ed2a70148, plus the later default-on flip from 573d44dacf) -- only the linear half of that work had made it into release14-quant-shaders so far (f60080d04a), leaving SDPA on the tiled path even with ET_VK_TEXTURE_COOPMAT=1 set. Same gating as dev: enabled by default on capability-eligible devices (cooperative_matrix + subgroup_size==64), ET_VK_DISABLE_COOPMAT=1 remains the shared kill switch. No texture3d IO work was needed -- SDPA's q/k/v/attn_weights tensors are already buffer+half in the real e2e graph, which is what the existing buffer-storage-only coopmat shaders require. Validated via ETDump on primary M51: sdpa_compute_attn_weights_coopmat and sdpa_compute_out_coopmat now dispatch (previously silently fell back to tiled). 1B/4w prefill 1075 -> 1508 tok/s stacking this on top of the existing linear WMMA. Authored with Claude Code. --- .../sdpa_compute_attn_weights_coopmat.glsl | 277 ++++++++++++++++++ .../sdpa_compute_attn_weights_coopmat.yaml | 37 +++ .../ops/glsl/sdpa_compute_out_coopmat.glsl | 245 ++++++++++++++++ .../ops/glsl/sdpa_compute_out_coopmat.yaml | 33 +++ .../vulkan/runtime/graph/ops/impl/SDPA.cpp | 162 +++++++++- 5 files changed, 748 insertions(+), 6 deletions(-) create mode 100644 backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_coopmat.glsl create mode 100644 backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_coopmat.yaml create mode 100644 backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coopmat.glsl create mode 100644 backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coopmat.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_coopmat.glsl b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_coopmat.glsl new file mode 100644 index 00000000000..ac24c1d8d59 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_coopmat.glsl @@ -0,0 +1,277 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +/* + * KHR Cooperative Matrix SDPA QK^T kernel (prefill / LLM mode). + * + * Computes per head q_h: attn[s, c] = inv_scale * sum_d Q[s, d] * K[c, d] + * Q = q_projected, DHSB [S, Q_H, D] index (vec4 along d): (s*Q_H + q_h)*D4 + d4 + * K = k_cache, DHSB [context_len, KV_H, D] index: (c*KV_H + kv_h)*D4 + d4 + * (GQA: kv_h = q_h / (Q_H/KV_H)) + * attn = attn_weights, head-contiguous [S_aligned, C4*4] + * scalar index: (q_h*S_aligned + s)*(C4*4) + c + * then the causal mask sets attn[s,c] = -inf where c > s + input_pos. + * + * Reduction dim is D (head_dim) -> num_k_chunks = D / WG_TILE_K (2 or 4). + * + * Two structural differences from coopmat_mm.glsl: + * 1. K is consumed transposed (we need Q*K^T). Rather than a ColumnMajor load + * of packed shared memory, K is staged TRANSPOSED into an fp16 shared array + * laid out [d][c] (scatter on write, since native K has d contiguous), so + * the MMA loop reads it RowMajor exactly like A. Q is likewise staged into + * an fp16 [s][d] shared array. + * 2. The causal mask cannot be applied to a coopmat accumulator (opaque + * lane->element mapping), so the scaled fp16 result is coopMatStore'd to a + * shared [s][c] scratch and then copied to global scalar-wise, applying the + * per-element mask. A whole-WG-tile that is entirely above the diagonal is + * written as -inf and skips the MMA loop (~halves prefill QK^T work). + * + * Dispatch: global {num_tiles_n*WG_SIZE, num_tiles_m, H_q}, local {WG_SIZE,1,1}. + * tileID = gl_WorkGroupID.xy (x->context, y->seq), q_h = gl_WorkGroupID.z. + */ + +#version 450 core + +#extension GL_KHR_cooperative_matrix : require +#extension GL_KHR_memory_scope_semantics : require +#extension GL_KHR_shader_subgroup_basic : enable +#extension GL_EXT_shader_explicit_arithmetic_types : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_EXT_control_flow_attributes : enable + +#define PRECISION ${PRECISION} + +layout(std430) buffer; + +#include "common.glslh" + +// Bindings mirror sdpa_compute_attn_weights_tiled: attn_weights(0), q(1), k(2). +// attn_weights is written scalar-wise (masked copy), so declare scalar array. +// Coopmat is buffer-only; IO_STORAGE / K_CACHE_STORAGE are always buffer here +// (the yaml only generates the buffer/buffer variant) but are kept as params so +// the generated name carries the _buffer_buffer suffix the dispatch builds. +${layout_declare_tensor(B, "w", "t_attn_weights", DTYPE, IO_STORAGE, is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_q", DTYPE, IO_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_k", DTYPE, K_CACHE_STORAGE, is_scalar_array=False)} + +${layout_declare_ubo(B, "ivec4", "q_sizes")} +${layout_declare_ubo(B, "ivec4", "k_sizes")} +${layout_declare_ubo(B, "int", "input_pos")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +${layout_declare_spec_const(C, "float", "inv_scale", "1.0")} +// K-chunk trip count = head_dim / WG_TILE_K, as a spec constant (the +// Xclipse/AMD-PAL compiler crashes on a coopMatMulAdd loop with a UBO-derived +// trip count — see coopmat_mm.glsl). +${layout_declare_spec_const(C, "int", "num_k_chunks_arg", "0")} + +const uint MMA_M = ${MMA_M}; +const uint MMA_N = ${MMA_N}; +const uint MMA_K = ${MMA_K}; + +const uint WG_TILE_M = ${WG_TILE_M}; +const uint WG_TILE_N = ${WG_TILE_N}; +const uint WG_TILE_K = ${WG_TILE_K}; + +const uint SG_GRID_X = ${SG_GRID_X}; +const uint SG_GRID_Y = ${SG_GRID_Y}; +const uint SUBGROUP_SIZE = ${SUBGROUP_SIZE}; +const uint NUM_SUBGROUPS = SG_GRID_X * SG_GRID_Y; +const uint WG_SIZE = NUM_SUBGROUPS * SUBGROUP_SIZE; + +const uint SG_TILE_M = WG_TILE_M / SG_GRID_Y; +const uint SG_TILE_N = WG_TILE_N / SG_GRID_X; +const uint MMAS_PER_SG_M = SG_TILE_M / MMA_M; +const uint MMAS_PER_SG_N = SG_TILE_N / MMA_N; + +const uint FP16_PER_VEC4 = 4; // we read native tensors as f16vec4 (4 fp16) + +// fp16 shared tiles with skew padding. A = Q [s][d], B = K^T [d][c], +// C = scaled result [s][c] scratch for the masked scalar store. +const uint A_PAD = 8; +const uint B_PAD = 8; +const uint A_ROW = WG_TILE_K + A_PAD; +const uint B_ROW = WG_TILE_N + B_PAD; + +shared float16_t Ash[WG_TILE_M * A_ROW]; +shared float16_t Bsh[WG_TILE_K * B_ROW]; +shared float16_t Csh[WG_TILE_M * WG_TILE_N]; + +coopmat result[MMAS_PER_SG_M][MMAS_PER_SG_N]; + +void main() { + const uvec2 tileID = uvec2(gl_WorkGroupID.xy); + const uvec2 warpInTile = uvec2( + gl_SubgroupID % SG_GRID_X, + gl_SubgroupID / SG_GRID_X); + const int q_h = int(gl_WorkGroupID.z); + + // LLM layout: q_sizes WHCN {D, H_q, S, B}; k_sizes WHCN {D, H_kv, C_max, B}. + const int D = q_sizes.x; + const int Q_H = q_sizes.y; + const int S = q_sizes.z; + const int KV_H = k_sizes.y; + const int D4 = div_up_4(D); + const int S_aligned = align_up_4(S); + const int context_len = input_pos + S; + const int C4 = div_up_4(context_len); + const int aw_row_width = C4 * 4; + + int kv_h = q_h; + if (KV_H < Q_H) { + kv_h = q_h / (Q_H / KV_H); + } + + const uint M = uint(S); // output rows (seq) + const uint N = uint(context_len); // output cols (context) + const uint num_tiles_n = (N + WG_TILE_N - 1u) / WG_TILE_N; + const uint num_tiles_m = (M + WG_TILE_M - 1u) / WG_TILE_M; + if (tileID.x >= num_tiles_n || tileID.y >= num_tiles_m) { + return; + } + + const uint s_tile_base = WG_TILE_M * tileID.y; + const uint c_tile_base = WG_TILE_N * tileID.x; + + const float16_t NEG_INF = float16_t(-1.0 / 0.0); + + // Whole-tile causal skip: if the lowest context index in this tile exceeds + // the highest (s + input_pos), every element is masked. + const bool tile_all_masked = + int(c_tile_base) > (int(s_tile_base) + int(WG_TILE_M) - 1 + input_pos); + if (tile_all_masked) { + for (uint idx = gl_LocalInvocationID.x; idx < WG_TILE_M * WG_TILE_N; + idx += WG_SIZE) { + const uint ls = idx / WG_TILE_N; + const uint lc = idx % WG_TILE_N; + const uint gs = s_tile_base + ls; + const uint gc = c_tile_base + lc; + if (gs < uint(S) && gc < uint(context_len)) { + t_attn_weights[(uint(q_h) * uint(S_aligned) + gs) * + uint(aw_row_width) + + gc] = NEG_INF; + } + } + return; + } + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + result[i][j] = coopmat(0.0); + } + } + + // Per-head bases / row strides (vec4 units) for the native DHSB reads. + const uint q_head_base = uint(q_h) * uint(D4); + const uint q_row_stride = uint(Q_H) * uint(D4); + const uint k_head_base = uint(kv_h) * uint(D4); + const uint k_row_stride = uint(KV_H) * uint(D4); + + const uint VEC4_PER_CHUNK = WG_TILE_K / 4u; // d4 columns per K-chunk + + for (uint chunk = 0; chunk < uint(num_k_chunks_arg); ++chunk) { + const uint d4_chunk = chunk * VEC4_PER_CHUNK; // d4 offset of this chunk + + // --- Stage A = Q [s][d] into fp16 shared (contiguous d) --- + for (uint idx = gl_LocalInvocationID.x; idx < WG_TILE_M * VEC4_PER_CHUNK; + idx += WG_SIZE) { + const uint ls = idx / VEC4_PER_CHUNK; // local s row + const uint ld4 = idx % VEC4_PER_CHUNK; // local d4 within chunk + const uint gs = s_tile_base + ls; + f16vec4 v = t_q[gs * q_row_stride + q_head_base + d4_chunk + ld4]; + const uint base = ls * A_ROW + ld4 * 4u; + Ash[base + 0u] = v.x; + Ash[base + 1u] = v.y; + Ash[base + 2u] = v.z; + Ash[base + 3u] = v.w; + } + + // --- Stage B = K^T [d][c] into fp16 shared (transpose on write) --- + for (uint idx = gl_LocalInvocationID.x; idx < WG_TILE_N * VEC4_PER_CHUNK; + idx += WG_SIZE) { + const uint lc = idx / VEC4_PER_CHUNK; // local c row of K + const uint ld4 = idx % VEC4_PER_CHUNK; // local d4 within chunk + const uint gc = c_tile_base + lc; + f16vec4 v = t_k[gc * k_row_stride + k_head_base + d4_chunk + ld4]; + const uint d_base = ld4 * 4u; // local d within chunk + Bsh[(d_base + 0u) * B_ROW + lc] = v.x; + Bsh[(d_base + 1u) * B_ROW + lc] = v.y; + Bsh[(d_base + 2u) * B_ROW + lc] = v.z; + Bsh[(d_base + 3u) * B_ROW + lc] = v.w; + } + + barrier(); + + // --- Cooperative matrix MMA: result += A * B (B is already K^T) --- + [[unroll]] for (uint k = 0; k < WG_TILE_K / MMA_K; ++k) { + uint k_start = MMA_K * k; + + coopmat matA[MMAS_PER_SG_M]; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + uint row_a = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + coopMatLoad( + matA[i], Ash, + row_a * A_ROW + k_start, + A_ROW, + gl_CooperativeMatrixLayoutRowMajor); + } + + coopmat matB; + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + uint col_b = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopMatLoad( + matB, Bsh, + k_start * B_ROW + col_b, + B_ROW, + gl_CooperativeMatrixLayoutRowMajor); + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + result[i][j] = coopMatMulAdd(matA[i], matB, result[i][j]); + } + } + } + + barrier(); + } + + // --- Scale on the fp32 accumulator, store fp16 into Csh [s][c] scratch --- + const float16_t inv_scale_h = float16_t(inv_scale); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + result[i][j] = result[i][j] * inv_scale; // fp32 scalar multiply + coopmat out_tile = + coopmat(result[i][j]); + uint local_row = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + uint local_col = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopMatStore( + out_tile, Csh, + local_row * WG_TILE_N + local_col, WG_TILE_N, + gl_CooperativeMatrixLayoutRowMajor); + } + } + barrier(); + + // --- Copy Csh -> global attn_weights with the per-element causal mask --- + for (uint idx = gl_LocalInvocationID.x; idx < WG_TILE_M * WG_TILE_N; + idx += WG_SIZE) { + const uint ls = idx / WG_TILE_N; + const uint lc = idx % WG_TILE_N; + const uint gs = s_tile_base + ls; + const uint gc = c_tile_base + lc; + if (gs < uint(S) && gc < uint(context_len)) { + float16_t v = Csh[idx]; + if (int(gc) > int(gs) + input_pos) { + v = NEG_INF; + } + t_attn_weights[(uint(q_h) * uint(S_aligned) + gs) * + uint(aw_row_width) + + gc] = v; + } + } +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_coopmat.yaml b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_coopmat.yaml new file mode 100644 index 00000000000..26b516df732 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_coopmat.yaml @@ -0,0 +1,37 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# KHR Cooperative Matrix SDPA QK^T kernel (prefill / LLM). Buffer-only, fp16. +# Variant name matches the dispatch: +# sdpa_compute_attn_weights_coopmat_buffer_buffer_half. + +sdpa_compute_attn_weights_coopmat: + parameter_names_with_default_values: + DTYPE: half + PRECISION: highp + IO_STORAGE: buffer + K_CACHE_STORAGE: buffer + MMA_M: 16 + MMA_N: 16 + MMA_K: 16 + # 128x64 is the tile-sweep optimum that FITS this shader's shared memory + # (the generic-matmul sweep's 128x128 needs ~50KB LDS once the causal-mask + # Csh scratch is added — overflows M5 EVT1). M-tile 128, N-tile 64. + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + generate_variant_forall: + combination: + parameter_names: [IO_STORAGE, K_CACHE_STORAGE] + combos: + - parameter_values: [buffer, buffer] + DTYPE: + - VALUE: half + shader_variants: + - NAME: sdpa_compute_attn_weights_coopmat diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coopmat.glsl b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coopmat.glsl new file mode 100644 index 00000000000..d3173441ccb --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coopmat.glsl @@ -0,0 +1,245 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +/* + * KHR Cooperative Matrix SDPA attn*V kernel (prefill / LLM mode). + * + * Computes per head q_h: out[s, d] = sum_c P[s, c] * V[c, d] + * P = attn_weights (softmax output), head-contiguous [S_aligned, context_len] + * index (vec4 along c): (q_h * S_aligned + s) * C4 + c4 + * V = v_cache, DHSB [context_len, KV_H, D] + * index (vec4 along d): (c * KV_H + kv_h) * D4 + d4 (GQA: kv_h = q_h/(Q_H/KV_H)) + * out= DHSB [S, Q_H, D] + * scalar index: (s * Q_H + q_h) * D + d + * + * This is the plain A*B coopmat MM (coopmat_mm.glsl) with three changes: + * - A (P) staging uses the head-contiguous row stride C4 + per-head base. + * - B (V) staging uses the DHSB head-interleaved row stride KV_H*D4 + base. + * - output coopMatStore uses the DHSB row stride (Q_H*D) so heads interleave; + * stride + head_dim are spec constants (the Xclipse/AMD-PAL compiler + * miscompiles coopMatStore whose stride derives from a UBO value). + * fp16 x fp16 -> fp32 MMA. No mask / no scale (softmax already applied). + * + * Dispatch: global {num_tiles_n*WG_SIZE, num_tiles_m, H_q}, local {WG_SIZE,1,1}. + * tileID = gl_WorkGroupID.xy, q_h = gl_WorkGroupID.z. + */ + +#version 450 core + +#extension GL_KHR_cooperative_matrix : require +#extension GL_KHR_memory_scope_semantics : require +#extension GL_KHR_shader_subgroup_basic : enable +#extension GL_EXT_shader_explicit_arithmetic_types : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_EXT_control_flow_attributes : enable + +#define PRECISION ${PRECISION} + +layout(std430) buffer; + +#include "common.glslh" + +// Bindings mirror sdpa_compute_out_tiled: output(0), attn_weights(1), v(2). +// Coopmat is buffer-only; IO_STORAGE / V_CACHE_STORAGE are always buffer here +// (the yaml only generates the buffer/buffer variant) but are kept as params so +// the generated name carries the _buffer_buffer suffix the dispatch builds. +${layout_declare_tensor(B, "w", "t_output", DTYPE, IO_STORAGE, is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_attn_weights", DTYPE, IO_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_v", DTYPE, V_CACHE_STORAGE, is_scalar_array=False)} + +${layout_declare_ubo(B, "ivec4", "q_sizes")} +${layout_declare_ubo(B, "ivec4", "v_sizes")} +${layout_declare_ubo(B, "int", "input_pos")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +// Spec constants. inv_scale occupies id 3 (UNUSED here — softmax already +// normalized; attn*V applies no scale) so this shader stays aligned with the +// decode _coop / tiled attn*V variants that share this node's fixed spec_vars +// list and declare inv_scale at id 3. The rest are never UBO-derived (see +// coopmat_mm.glsl for the Xclipse PAL bug these work around): K-chunk trip count +// (= max_context_len/WG_TILE_K), the DHSB output row stride (Q_H*D), and +// head_dim (D) for the store column offset. +${layout_declare_spec_const(C, "float", "inv_scale_unused", "1.0")} +${layout_declare_spec_const(C, "int", "num_k_chunks_arg", "0")} +${layout_declare_spec_const(C, "int", "out_row_stride_arg", "0")} +${layout_declare_spec_const(C, "int", "head_dim_arg", "0")} + +const uint MMA_M = ${MMA_M}; +const uint MMA_N = ${MMA_N}; +const uint MMA_K = ${MMA_K}; + +const uint WG_TILE_M = ${WG_TILE_M}; +const uint WG_TILE_N = ${WG_TILE_N}; +const uint WG_TILE_K = ${WG_TILE_K}; + +const uint SG_GRID_X = ${SG_GRID_X}; +const uint SG_GRID_Y = ${SG_GRID_Y}; +const uint SUBGROUP_SIZE = ${SUBGROUP_SIZE}; +const uint NUM_SUBGROUPS = SG_GRID_X * SG_GRID_Y; +const uint WG_SIZE = NUM_SUBGROUPS * SUBGROUP_SIZE; + +const uint SG_TILE_M = WG_TILE_M / SG_GRID_Y; +const uint SG_TILE_N = WG_TILE_N / SG_GRID_X; +const uint MMAS_PER_SG_M = SG_TILE_M / MMA_M; +const uint MMAS_PER_SG_N = SG_TILE_N / MMA_N; + +const uint FP16_PER_VEC4 = 8; + +const uint A_STRIDE_VEC4 = (WG_TILE_K + FP16_PER_VEC4) / FP16_PER_VEC4; +const uint B_STRIDE_VEC4 = (WG_TILE_N + FP16_PER_VEC4) / FP16_PER_VEC4; + +shared uvec4 Ash[WG_TILE_M * A_STRIDE_VEC4]; +shared uvec4 Bsh[WG_TILE_K * B_STRIDE_VEC4]; + +coopmat result[MMAS_PER_SG_M][MMAS_PER_SG_N]; + +void main() { + const uvec2 tileID = uvec2(gl_WorkGroupID.xy); + const uvec2 warpInTile = uvec2( + gl_SubgroupID % SG_GRID_X, + gl_SubgroupID / SG_GRID_X); + const int q_h = int(gl_WorkGroupID.z); + + // LLM layout: q_sizes WHCN {D, H_q, S, B}; v_sizes WHCN {D, H_kv, C_max, B}. + const int D = q_sizes.x; + const int Q_H = q_sizes.y; + const int S = q_sizes.z; + const int KV_H = v_sizes.y; + const int D4 = div_up_4(D); + const int S_aligned = align_up_4(S); + const int context_len = input_pos + S; + const int C4 = div_up_4(context_len); + + int kv_h = q_h; + if (KV_H < Q_H) { + kv_h = q_h / (Q_H / KV_H); + } + + const uint M = uint(S); // output rows + const uint N = uint(head_dim_arg); // output cols (= D) + const uint num_tiles_n = (N + WG_TILE_N - 1u) / WG_TILE_N; + const uint num_tiles_m = (M + WG_TILE_M - 1u) / WG_TILE_M; + if (tileID.x >= num_tiles_n || tileID.y >= num_tiles_m) { + return; + } + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + result[i][j] = coopmat(0.0); + } + } + + const uint INVS_PER_ROW_A = WG_TILE_K / FP16_PER_VEC4; + const uint a_col = gl_LocalInvocationID.x % INVS_PER_ROW_A; + const uint a_row_offset = gl_LocalInvocationID.x / INVS_PER_ROW_A; + + const uint INVS_PER_ROW_B = WG_TILE_N / FP16_PER_VEC4; + const uint b_col = gl_LocalInvocationID.x % INVS_PER_ROW_B; + const uint b_row_offset = gl_LocalInvocationID.x / INVS_PER_ROW_B; + + const uint a_row_base = WG_TILE_M * tileID.y; // global s + const uint b_col_base = WG_TILE_N * tileID.x; // global d + + // Per-head bases / strides in vec4 units. + const uint aw_head_base = uint(q_h) * uint(S_aligned) * uint(C4); + const uint aw_row_stride = uint(C4); + const uint v_head_base = uint(kv_h) * uint(D4); + const uint v_row_stride = uint(KV_H) * uint(D4); + + for (uint chunk = 0; chunk < uint(num_k_chunks_arg); ++chunk) { + const uint chunkK = chunk * WG_TILE_K; // along context_len + // num_k_chunks is max_context_len/WG_TILE_K (static spec const). The + // gate guarantees context_len % WG_TILE_N == 0, hence % WG_TILE_K == 0, + // so a chunk is either fully within context_len or fully beyond it. + // Stage zeros for beyond-context chunks (zero contribution to the MMA). + const bool chunk_valid = chunkK < uint(context_len); + + // --- Load A (attn_weights) tile -> shared (single pass) --- + { + f16vec4 v0 = f16vec4(0); + f16vec4 v1 = f16vec4(0); + if (chunk_valid) { + uint row = a_row_base + a_row_offset; // global s + uint k_hv4 = (chunkK + a_col * FP16_PER_VEC4) / 4u; // c, vec4 + uint base = aw_head_base + row * aw_row_stride; + v0 = t_attn_weights[base + k_hv4]; + v1 = t_attn_weights[base + k_hv4 + 1u]; + } + Ash[a_row_offset * A_STRIDE_VEC4 + a_col] = uvec4( + packFloat2x16(v0.xy), packFloat2x16(v0.zw), + packFloat2x16(v1.xy), packFloat2x16(v1.zw)); + } + + // --- Load B (V) tile -> shared (single pass), row-major [c, d] --- + { + f16vec4 v0 = f16vec4(0); + f16vec4 v1 = f16vec4(0); + if (chunk_valid) { + uint k_row = chunkK + b_row_offset; // global c + uint n_elem = b_col_base + b_col * FP16_PER_VEC4; // d + uint n4_0 = n_elem >> 2u; // d4 + uint base = k_row * v_row_stride + v_head_base; + v0 = t_v[base + n4_0]; + v1 = t_v[base + n4_0 + 1u]; + } + Bsh[b_row_offset * B_STRIDE_VEC4 + b_col] = uvec4( + packFloat2x16(v0.xy), packFloat2x16(v0.zw), + packFloat2x16(v1.xy), packFloat2x16(v1.zw)); + } + + barrier(); + + // --- Cooperative matrix MMA (identical to coopmat_mm) --- + [[unroll]] for (uint k = 0; k < WG_TILE_K / MMA_K; ++k) { + uint k_start = MMA_K * k; + + coopmat matA[MMAS_PER_SG_M]; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + uint row_a = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + coopMatLoad( + matA[i], Ash, + row_a * A_STRIDE_VEC4 + k_start / FP16_PER_VEC4, + A_STRIDE_VEC4, + gl_CooperativeMatrixLayoutRowMajor); + } + + coopmat matB; + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + uint col_b = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j) / FP16_PER_VEC4; + coopMatLoad( + matB, Bsh, + k_start * B_STRIDE_VEC4 + col_b, + B_STRIDE_VEC4, + gl_CooperativeMatrixLayoutRowMajor); + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + result[i][j] = coopMatMulAdd(matA[i], matB, result[i][j]); + } + } + } + + barrier(); + } + + // --- Store result: DHSB out, head-interleaved row stride (spec const) --- + const uint out_row_stride = uint(out_row_stride_arg); // Q_H * D + const uint head_off = uint(q_h) * uint(head_dim_arg); // q_h * D + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + uint gi = WG_TILE_M * tileID.y + MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + uint gj = WG_TILE_N * tileID.x + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopmat out_tile = + coopmat(result[i][j]); + coopMatStore( + out_tile, t_output, + gi * out_row_stride + head_off + gj, out_row_stride, + gl_CooperativeMatrixLayoutRowMajor); + } + } +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coopmat.yaml b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coopmat.yaml new file mode 100644 index 00000000000..e5576003be7 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coopmat.yaml @@ -0,0 +1,33 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# KHR Cooperative Matrix SDPA attn*V kernel (prefill / LLM). Buffer-only, fp16. +# Variant name matches the dispatch: sdpa_compute_out_coopmat_buffer_buffer_half. + +sdpa_compute_out_coopmat: + parameter_names_with_default_values: + DTYPE: half + PRECISION: highp + IO_STORAGE: buffer + V_CACHE_STORAGE: buffer + MMA_M: 16 + MMA_N: 16 + MMA_K: 16 + WG_TILE_M: 64 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + generate_variant_forall: + combination: + parameter_names: [IO_STORAGE, V_CACHE_STORAGE] + combos: + - parameter_values: [buffer, buffer] + DTYPE: + - VALUE: half + shader_variants: + - NAME: sdpa_compute_out_coopmat diff --git a/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp b/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp index 3efb834725d..29469cc0770 100644 --- a/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp @@ -22,6 +22,7 @@ #include #include +#include namespace vkcompute { @@ -178,16 +179,90 @@ static inline SDPAMode mode_of(const std::vector& resize_args) { return static_cast(resize_args.at(3)); } +// +// Cooperative-matrix (WMMA) SDPA prefill path. +// +// The QK^T and attn*V coopmat shaders use a 64x64x32 WG tile with 4 subgroups +// of 64 lanes (256 invocations), matching coopmat_mm.glsl. Selected only for +// LLM prefill (S > 1) on a coopmat-capable discrete RDNA GPU with buffer/fp16 +// tensors and tile-aligned shapes; decode (S == 1) stays on the _coop GEMV +// path and ineligible shapes fall back to _tiled. Enabled by default on +// capability-eligible devices; ET_VK_DISABLE_COOPMAT remains the kill switch +// (shared with the q4gsw linear coopmat path). +// +constexpr uint32_t kSdpaCmTileM = 64; +constexpr uint32_t kSdpaCmTileN = 64; +constexpr uint32_t kSdpaCmTileK = 32; +constexpr uint32_t kSdpaCmInvocations = 256; +// QK^T uses a 128-tall M-tile (the 128x64 tile-sweep optimum that still fits +// the masked shader's shared memory); attn*V keeps the 64x64 geometry. Same +// WG_SIZE (256 = 2x2 subgroups x 64), so only the M-tile count differs for +// QK^T. +constexpr uint32_t kSdpaCmQkTileM = 128; + +static bool sdpa_coopmat_not_disabled() { + return std::getenv("ET_VK_DISABLE_COOPMAT") == nullptr; +} + +static bool sdpa_coopmat_device_ok(ComputeGraph* graph) { + if (!sdpa_coopmat_not_disabled()) { + return false; + } + const auto* adapter = graph->context()->adapter_ptr(); + // NOTE: intentionally NO !is_integrated_gpu() check. The target M5 EVT1 + // (Xclipse 970) is a unified-memory "integrated" GPU but has fast fp16 WMMA; + // the q4gsw linear coopmat gate likewise omits this check and runs there. The + // generic matmul gate keeps it (to avoid coopmat on iGPUs without WMMA); SDPA + // coopmat is enabled by default on capability-eligible devices, so the + // subgroup/cooperative-matrix checks below are the only gate. + return adapter->supports_cooperative_matrix() && + adapter->subgroup_size() == 64; +} + +static bool sdpa_buf_half(ComputeGraph* graph, const ValueRef t) { + return graph->storage_type_of(t) == utils::kBuffer && + graph->dtype_of(t) == vkapi::kHalf; +} + +static inline bool sdpa_cm_aligned(int64_t m, int64_t n, int64_t k) { + return m % static_cast(kSdpaCmTileM) == 0 && + n % static_cast(kSdpaCmTileN) == 0 && + k % static_cast(kSdpaCmTileK) == 0; +} + +static inline bool is_sdpa_coopmat(const vkapi::ShaderInfo& shader) { + return shader.kernel_name.find("_coopmat") != std::string::npos; +} + vkapi::ShaderInfo pick_sdpa_qk_shader( ComputeGraph* graph, const std::vector& args, const std::vector& resize_args) { const SDPAMode mode = mode_of(resize_args); if (mode == SDPAMode::LLM) { + const ValueRef attn_weights = args.at(0).refs.at(0); const ValueRef q_projected = args.at(1).refs.at(0); const ValueRef k_cache = args.at(1).refs.at(1); const bool is_gemv = is_single_token(graph, q_projected); + // Prefill WMMA path: Q @ K^T with K = head_dim, N = context_len, M = S. + if (!is_gemv && sdpa_coopmat_device_ok(graph) && + sdpa_buf_half(graph, q_projected) && sdpa_buf_half(graph, k_cache) && + sdpa_buf_half(graph, attn_weights)) { + const SDPADims d = compute_sdpa_dims( + *graph, q_projected, k_cache, resize_args.at(2), SDPAMode::LLM); + if (d.S % static_cast(kSdpaCmQkTileM) == 0 && + d.context_len % static_cast(kSdpaCmTileN) == 0 && + d.D % static_cast(kSdpaCmTileK) == 0) { + std::string shader_name = "sdpa_compute_attn_weights_coopmat"; + add_storage_type_suffix( + shader_name, graph->storage_type_of(q_projected)); + add_storage_type_suffix(shader_name, graph->storage_type_of(k_cache)); + add_dtype_suffix(shader_name, graph->dtype_of(q_projected)); + return VK_KERNEL_FROM_STR(shader_name); + } + } + std::string shader_name = "sdpa_compute_attn_weights"; shader_name += is_gemv ? "_coop" : "_tiled"; add_storage_type_suffix(shader_name, graph->storage_type_of(q_projected)); @@ -214,7 +289,6 @@ utils::uvec3 pick_sdpa_qk_global_wg_size( const vkapi::ShaderInfo& shader, const std::vector& args, const std::vector& resize_args) { - (void)shader; (void)args; const SDPAMode mode = mode_of(resize_args); const ValueRef q = resize_args.at(0); @@ -222,6 +296,20 @@ utils::uvec3 pick_sdpa_qk_global_wg_size( const ValueRef input_pos_symint = resize_args.at(2); const SDPADims d = compute_sdpa_dims(*graph, q, k, input_pos_symint, mode); + if (is_sdpa_coopmat(shader)) { + // One workgroup per 64x64 output tile (N = context_len, M = S); + // *kSdpaCmInvocations cancels the framework div_up against local x. z + // carries the head index. + const uint32_t num_tiles_n = + utils::div_up(static_cast(d.context_len), kSdpaCmTileN); + const uint32_t num_tiles_m = + utils::div_up(static_cast(d.S), kSdpaCmQkTileM); + return { + num_tiles_n * kSdpaCmInvocations, + num_tiles_m, + static_cast(d.H * d.B)}; + } + // Dispatch grid: (context_len tiles, S tiles, H * B). const uint32_t N4 = utils::div_up_4(static_cast(d.context_len)); const uint32_t M4 = utils::div_up_4(static_cast(d.S)); @@ -236,6 +324,11 @@ utils::uvec3 pick_sdpa_qk_local_wg_size( const std::vector& resize_args) { const SDPAMode mode = mode_of(resize_args); if (mode == SDPAMode::LLM) { + // _coopmat must be checked before _coop (the former contains the latter as + // a substring); the coopmat shaders use a flat 256-lane workgroup. + if (is_sdpa_coopmat(shader)) { + return {kSdpaCmInvocations, 1, 1}; + } const bool use_coop_algorithm = shader.kernel_name.find("_coop") != std::string::npos; if (use_coop_algorithm) { @@ -293,10 +386,31 @@ vkapi::ShaderInfo pick_sdpa_av_shader( const SDPAMode mode = mode_of(resize_args); if (mode == SDPAMode::LLM) { const ValueRef out = args.at(0).refs.at(0); + const ValueRef attn_weights_softmax = args.at(1).refs.at(0); const ValueRef v_cache = args.at(1).refs.at(1); const ValueRef q_projected = resize_args.at(0); const bool is_gemv = is_single_token(graph, q_projected); + // Prefill WMMA path: P @ V with K = context_len, N = head_dim, M = S. + if (!is_gemv && sdpa_coopmat_device_ok(graph) && + sdpa_buf_half(graph, out) && + sdpa_buf_half(graph, attn_weights_softmax) && + sdpa_buf_half(graph, v_cache)) { + const SDPADims d = compute_sdpa_dims( + *graph, + q_projected, + resize_args.at(1), + resize_args.at(2), + SDPAMode::LLM); + if (sdpa_cm_aligned(/*m=*/d.S, /*n=*/d.D, /*k=*/d.context_len)) { + std::string shader_name = "sdpa_compute_out_coopmat"; + add_storage_type_suffix(shader_name, graph->storage_type_of(out)); + add_storage_type_suffix(shader_name, graph->storage_type_of(v_cache)); + add_dtype_suffix(shader_name, graph->dtype_of(out)); + return VK_KERNEL_FROM_STR(shader_name); + } + } + std::string shader_name = "sdpa_compute_out"; shader_name += is_gemv ? "_coop" : "_tiled"; add_storage_type_suffix(shader_name, graph->storage_type_of(out)); @@ -319,13 +433,24 @@ utils::uvec3 pick_sdpa_av_global_wg_size( const vkapi::ShaderInfo& shader, const std::vector& args, const std::vector& resize_args) { - (void)shader; const SDPAMode mode = mode_of(resize_args); const ValueRef q = resize_args.at(0); const ValueRef k = resize_args.at(1); const ValueRef input_pos_symint = resize_args.at(2); const SDPADims d = compute_sdpa_dims(*graph, q, k, input_pos_symint, mode); + if (is_sdpa_coopmat(shader)) { + // One workgroup per 64x64 output tile (N = head_dim, M = S). z = head. + const uint32_t num_tiles_n = + utils::div_up(static_cast(d.D), kSdpaCmTileN); + const uint32_t num_tiles_m = + utils::div_up(static_cast(d.S), kSdpaCmTileM); + return { + num_tiles_n * kSdpaCmInvocations, + num_tiles_m, + static_cast(d.H * d.B)}; + } + const uint32_t N4 = utils::div_up_4(static_cast(d.D)); const uint32_t M4 = utils::div_up_4(static_cast(d.S)); return {N4, M4, static_cast(d.H * d.B)}; @@ -339,6 +464,11 @@ utils::uvec3 pick_sdpa_av_local_wg_size( const std::vector& resize_args) { const SDPAMode mode = mode_of(resize_args); if (mode == SDPAMode::LLM) { + // _coopmat must be checked before _coop (the former contains the latter as + // a substring); the coopmat shaders use a flat 256-lane workgroup. + if (is_sdpa_coopmat(shader)) { + return {kSdpaCmInvocations, 1, 1}; + } const bool use_coop_algorithm = shader.kernel_name.find("_coop") != std::string::npos; if (use_coop_algorithm) { @@ -430,8 +560,12 @@ void add_sdpa_compute_attn_weights_node( param_ubos, // Push Constants {}, - // Specialization Constants - {scale_val}, + // Specialization Constants: {inv_scale (id 3), num_k_chunks (id 4)}. + // num_k_chunks = head_dim / WG_TILE_K is static and consumed only by the + // coopmat QK^T variant; the tiled/coop variants declare only id 3 and + // ignore the trailing entry. + {scale_val, + graph.size_at(-1, q) / static_cast(kSdpaCmTileK)}, // Resize Args: [q, k, input_pos_symint_or_dummy, mode] {q, k, input_pos_symint, mode_ref}, // Resizing Logic @@ -511,6 +645,20 @@ void add_sdpa_compute_out_node( const ValueRef mode_ref = static_cast(mode); + // Coopmat attn*V spec constants (static; consumed only by the coopmat + // variant — the tiled/coop variants ignore the trailing entries, and id 3 is + // the inv_scale slot the decode _coop shader reads, kept at 1.0 = no-op). + // num_k_chunks uses max_context_len (the loop bound is a spec const per the + // Xclipse bug); beyond-context chunks are zero-staged in the shader. + // Values are meaningful only in LLM mode; in FUSED they are ignored. + const int32_t cm_head_dim = graph.size_at(-1, q); + const int32_t cm_num_q_heads = graph.size_at(-2, q); + const int32_t cm_max_context = graph.size_at(-3, v); + const int32_t cm_num_k_chunks = + (cm_max_context + static_cast(kSdpaCmTileK) - 1) / + static_cast(kSdpaCmTileK); + const int32_t cm_out_row_stride = cm_num_q_heads * cm_head_dim; + graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, pick_sdpa_av_shader, @@ -522,8 +670,10 @@ void add_sdpa_compute_out_node( param_ubos, // Push Constants {}, - // Specialization Constants - {}, + // Specialization Constants: + // {inv_scale slot (id 3), num_k_chunks (id 4), out_row_stride (id 5), + // head_dim (id 6)}. + {1.0f, cm_num_k_chunks, cm_out_row_stride, cm_head_dim}, // Resize Args: [q, k, input_pos_symint_or_dummy, mode] {q, k, input_pos_symint, mode_ref}, // Resizing Logic From d7360b7a147b4665a95f3df1b4119c0a0b6297af Mon Sep 17 00:00:00 2001 From: Yanwen Xu Date: Tue, 18 Aug 2026 14:31:47 -0700 Subject: [PATCH 09/28] [ET-VK] Add SDPA coopmat correctness harness (--sdpa-correctness-only) run_sdpa_suite only measures timing for sdpa_compute_attn_weights_coopmat / sdpa_compute_out_coopmat -- it never reads output data back or checks it against a reference. Add a correctness gate: builds the same graph shape directly via ComputeGraph (SymInt support the TestCase framework lacks), at small coopmat-tile-aligned shapes (S=128, D=64, matching the 1B model's real Q_H/KV_H config as one case), computes a causal GQA-aware fp32 CPU reference matching the shader's exact head-mapping (kv_h = q_h / (Q_H/KV_H)), and confirms both coopmat shaders actually dispatched (not a silent tiled fallback, which would pass numerically too). Verified 10 back-to-back --sdpa-correctness-only runs, 0 failures across both cases (20/20) -- the repeat-run discipline the linear tile sweep's flaky-tile incident earlier today established as mandatory for any coopmat correctness claim, not just a single pass. --- .../test/custom_ops/test_llama_microbench.cpp | 264 +++++++++++++++++- 1 file changed, 263 insertions(+), 1 deletion(-) diff --git a/backends/vulkan/test/custom_ops/test_llama_microbench.cpp b/backends/vulkan/test/custom_ops/test_llama_microbench.cpp index 99693ebeac8..d566bc5eb2b 100644 --- a/backends/vulkan/test/custom_ops/test_llama_microbench.cpp +++ b/backends/vulkan/test/custom_ops/test_llama_microbench.cpp @@ -34,6 +34,15 @@ // Other flags: // --model= only run models whose name contains // --correctness-only run just the linear correctness matrix, skip perf +// --sdpa-correctness-only run just the SDPA coopmat correctness cases +// (sdpa_compute_attn_weights_coopmat / +// sdpa_compute_out_coopmat vs. a CPU causal-attention +// reference, at small tile-aligned shapes -- see +// run_sdpa_correctness). A single pass is NOT +// sufficient evidence for a coopmat shader (the linear +// tile sweep found a tile that passed once and then +// failed 1-in-10 identical repeats, silently) -- rerun +// this flag repeatedly before trusting a pass. // --skip-correctness skip the linear correctness gate before perf // --list print every case that would run (with sizes), no GPU // --help @@ -70,6 +79,7 @@ #include #include #include +#include #include #include #include @@ -1153,6 +1163,250 @@ bool run_sdpa_suite(const std::string& model_filter) { return all_confirmed; } +// ===================== sdpa correctness ===================== +// Coopmat SDPA correctness gate for sdpa_compute_attn_weights_coopmat +// (QK^T) and sdpa_compute_out_coopmat (attn*V). run_sdpa_suite above is +// perf-only -- it never reads output data back or checks it against a +// reference. Like the SDPA perf path (specs/010 Decision 8), this builds +// the graph directly via ComputeGraph rather than the TestCase framework, +// because llama.custom_sdpa.default needs SymInt support the framework +// doesn't have. +// +// Shapes are small and coopmat-tile-aligned (S%128==0, D%64==0 -- SDPA.cpp's +// kSdpaCmQkTileM/kSdpaCmTileN/kSdpaCmTileK/kSdpaCmTileM eligibility check) +// with input_pos=0, so context_len == seq_len and the freshly-written KV +// cache covers exactly the self-attention window under test -- no +// uninitialized cache history to account for. A naive O(S^2*D) fp32 +// reference at production S=2048 would be far too slow to run the repeated +// (10+) back-to-back passes a coopmat correctness check needs (today's +// linear-tile-sweep incident: a tile that passed --correctness-only ONCE +// was later found to fail 1-in-10 identical repeat runs, silent wrong +// output, no crash) -- these shapes keep one pass under a second. +struct SdpaCorrectnessCase { + const char* name; + int64_t seq_len; // S; also context_len (input_pos=0, fresh prefill) + int64_t head_dim; // D + int64_t num_heads; // Q_H + int64_t num_kv_heads; // KV_H +}; +const std::vector kSdpaCorrectnessCases = { + // Minimal GQA case: 128 is the smallest legal QK^T M-tile multiple, 64 + // the smallest legal head_dim (both QK^T K-tile and attn*V N-tile). + {"tiny_gqa", 128, 64, 2, 1}, + // 1B's real head configuration (head_dim=64, 32 Q heads, 8 KV heads -- + // kSdpaModels), S truncated from the real 2048 to the same aligned 128 + // so the CPU reference stays fast; this is the shape most likely to + // exercise a head-indexing (GQA) bug the tiny case's group size of 2 + // could hide. + {"1b_head_config", 128, 64, 32, 8}, +}; + +// Causal, GQA-aware fp32 CPU reference. q is [S, Q_H, D], k/v are +// [S, KV_H, D] (row-major, batch=1 squeezed). kv_h = q_h / (Q_H / KV_H), +// matching sdpa_compute_attn_weights_coopmat.glsl's GQA head mapping exactly +// (see that file's header comment). Causal: query s attends to context +// c <= s (input_pos=0). +std::vector sdpa_reference( + const std::vector& q, + const std::vector& k, + const std::vector& v, + int64_t S, + int64_t D, + int64_t Q_H, + int64_t KV_H) { + const float scale = 1.0f / std::sqrt(static_cast(D)); + const int64_t group = Q_H / KV_H; + std::vector out(static_cast(S * Q_H * D), 0.0f); + std::vector scores(static_cast(S)); + for (int64_t h = 0; h < Q_H; ++h) { + const int64_t kv_h = h / group; + for (int64_t s = 0; s < S; ++s) { + float max_score = -std::numeric_limits::infinity(); + for (int64_t c = 0; c <= s; ++c) { + float acc = 0.0f; + for (int64_t d = 0; d < D; ++d) { + acc += q[(s * Q_H + h) * D + d] * k[(c * KV_H + kv_h) * D + d]; + } + acc *= scale; + scores[c] = acc; + max_score = std::max(max_score, acc); + } + float denom = 0.0f; + for (int64_t c = 0; c <= s; ++c) { + scores[c] = std::exp(scores[c] - max_score); + denom += scores[c]; + } + for (int64_t d = 0; d < D; ++d) { + float acc = 0.0f; + for (int64_t c = 0; c <= s; ++c) { + acc += (scores[c] / denom) * v[(c * KV_H + kv_h) * D + d]; + } + out[(s * Q_H + h) * D + d] = acc; + } + } + } + return out; +} + +// Builds+runs one coopmat SDPA case via direct ComputeGraph construction +// (mirrors sdpa_run_case's graph shape), reads Q/K/V/out host-side in +// float/half, computes the CPU reference, and compares. Returns true iff +// BOTH the QK^T and attn*V coopmat shaders actually dispatched (not a +// silent tiled fallback -- tiled is also numerically correct, so a pure +// value comparison alone cannot tell the two apart) AND every output +// element is within tolerance. +bool sdpa_correctness_case(const SdpaCorrectnessCase& c) { + unsetenv("ET_VK_DISABLE_COOPMAT"); + + GraphConfig config; + config.enable_querypool = true; + api::context()->initialize_querypool(); + ComputeGraph graph(config); + + const int64_t B = 1; + const std::vector q_sizes = {B, c.seq_len, c.num_heads, c.head_dim}; + const std::vector kv_sizes = { + B, c.seq_len, c.num_kv_heads, c.head_dim}; + // context_len == seq_len: input_pos=0, so the cache is exactly this step's + // freshly-written K/V (see file comment above). + const std::vector cache_sizes = kv_sizes; + + IOValueRef r_q = + graph.add_input_tensor(q_sizes, vkapi::kHalf, utils::kBuffer); + IOValueRef r_k = + graph.add_input_tensor(kv_sizes, vkapi::kHalf, utils::kBuffer); + IOValueRef r_v = + graph.add_input_tensor(kv_sizes, vkapi::kHalf, utils::kBuffer); + + const ValueRef r_input_pos_symint = graph.add_symint(0); + const ValueRef r_out = + graph.add_tensor(q_sizes, vkapi::kHalf, utils::kBuffer); + + const ValueRef r_k_cache = + graph.add_tensor(cache_sizes, vkapi::kHalf, utils::kBuffer); + const ValueRef r_v_cache = + graph.add_tensor(cache_sizes, vkapi::kHalf, utils::kBuffer); + const ValueRef r_dummy_out = + graph.add_tensor({1}, vkapi::kHalf, utils::kBuffer); + + VK_GET_OP_FN("update_cache.default") + (graph, {r_k.value, r_k_cache, r_input_pos_symint, r_dummy_out}); + VK_GET_OP_FN("update_cache.default") + (graph, {r_v.value, r_v_cache, r_input_pos_symint, r_dummy_out}); + VK_GET_OP_FN("llama.custom_sdpa.default") + (graph, + { + r_q.value, + r_k_cache, + r_v_cache, + r_input_pos_symint, + kDummyValueRef, // attn_mask + kDummyValueRef, // dropout_p + kDummyValueRef, // is_causal + kDummyValueRef, // scale + r_out, + }); + + graph.set_output_tensor(r_out); + graph.prepare(); + graph.prepack(); + + const int64_t q_numel = c.seq_len * c.num_heads * c.head_dim; + const int64_t kv_numel = c.seq_len * c.num_kv_heads * c.head_dim; + + std::vector qf(q_numel), kf(kv_numel), vf(kv_numel); + std::vector qh(q_numel), kh(kv_numel), vh(kv_numel); + for (int64_t i = 0; i < q_numel; ++i) { + qf[i] = (static_cast(std::rand()) / RAND_MAX) * 2.0f - 1.0f; + qh[i] = float_to_half(qf[i]); + } + for (int64_t i = 0; i < kv_numel; ++i) { + kf[i] = (static_cast(std::rand()) / RAND_MAX) * 2.0f - 1.0f; + kh[i] = float_to_half(kf[i]); + vf[i] = (static_cast(std::rand()) / RAND_MAX) * 2.0f - 1.0f; + vh[i] = float_to_half(vf[i]); + } + graph.maybe_cast_and_copy_into_staging( + r_q.staging, qh.data(), static_cast(q_numel), vkapi::kHalf); + graph.maybe_cast_and_copy_into_staging( + r_k.staging, kh.data(), static_cast(kv_numel), vkapi::kHalf); + graph.maybe_cast_and_copy_into_staging( + r_v.staging, vh.data(), static_cast(kv_numel), vkapi::kHalf); + + graph.execute(); + + graph.context()->querypool().extract_results(); + const auto shader_results = + graph.context()->querypool().get_shader_timestamp_data(); + std::vector dispatched; + for (const auto& r : shader_results) { + dispatched.push_back(r.kernel_name); + } + const bool qk_fired = + has_kernel_containing(dispatched, "sdpa_compute_attn_weights_coopmat"); + const bool av_fired = + has_kernel_containing(dispatched, "sdpa_compute_out_coopmat"); + + std::vector outh(q_numel); + graph.maybe_cast_and_copy_from_staging( + graph.outputs()[0].staging, + outh.data(), + static_cast(q_numel), + vkapi::kHalf); + std::vector outf(q_numel); + for (int64_t i = 0; i < q_numel; ++i) { + outf[i] = half_to_float(outh[i]); + } + + const std::vector ref = sdpa_reference( + qf, kf, vf, c.seq_len, c.head_dim, c.num_heads, c.num_kv_heads); + + int64_t mismatches = 0; + int64_t first_mismatch = -1; + const float abs_tol = 0.03f; + const float rel_tol = 0.05f; + for (int64_t i = 0; i < q_numel; ++i) { + const float diff = std::fabs(outf[i] - ref[i]); + const float thresh = abs_tol + rel_tol * std::fabs(ref[i]); + if (diff > thresh) { + ++mismatches; + if (first_mismatch < 0) { + first_mismatch = i; + } + } + } + + const bool numeric_ok = mismatches == 0; + const bool fired_ok = qk_fired && av_fired; + std::cout << "[sdpa-correctness] " << c.name << " S=" << c.seq_len + << " D=" << c.head_dim << " Q_H=" << c.num_heads + << " KV_H=" << c.num_kv_heads + << " qk_coopmat=" << (qk_fired ? "yes" : "NO") + << " av_coopmat=" << (av_fired ? "yes" : "NO") + << " mismatches=" << mismatches << "/" << q_numel; + if (!numeric_ok) { + std::cout << " (first at " << first_mismatch << ": got=" << std::fixed + << std::setprecision(4) << outf[first_mismatch] + << " ref=" << ref[first_mismatch] << ")"; + } + std::cout << (numeric_ok && fired_ok ? " PASSED" : " FAILED") << "\n"; + return numeric_ok && fired_ok; +} + +// Runs every case in kSdpaCorrectnessCases once and reports pass/fail. +// Callers that want the repeated-pass discipline this coopmat bug class +// requires (see file comment above) should invoke this function itself +// multiple times in a loop -- kept a single pass per call, like +// run_linear_correctness, so a driver script controls the rep count and can +// distinguish "which specific repeat failed." +bool run_sdpa_correctness() { + bool all_ok = true; + for (const auto& c : kSdpaCorrectnessCases) { + all_ok = sdpa_correctness_case(c) && all_ok; + } + return all_ok; +} + // ============================== report ============================== const Record* find_record( @@ -1410,6 +1664,8 @@ void print_usage() { " correctness gate; a tile-variant token\n" " only affects one scheme+buffer cell)\n" " --correctness-only run just the linear correctness matrix\n" + " --sdpa-correctness-only run just the SDPA coopmat correctness " + "cases\n" " --skip-correctness skip the correctness gate before perf\n" " --list print every case with its sizes, no GPU\n" " --help this message\n"; @@ -1465,7 +1721,8 @@ void list_cases( int main(int argc, char** argv) { bool linear = false, baseline = false, sdpa = false; - bool correctness_only = false, skip_correctness = false, list_only = false; + bool correctness_only = false, sdpa_correctness_only = false; + bool skip_correctness = false, list_only = false; CaseFilter filter; for (int i = 1; i < argc; ++i) { const std::string arg = argv[i]; @@ -1477,6 +1734,8 @@ int main(int argc, char** argv) { sdpa = true; } else if (arg == "--correctness-only") { correctness_only = true; + } else if (arg == "--sdpa-correctness-only") { + sdpa_correctness_only = true; } else if (arg == "--skip-correctness") { skip_correctness = true; } else if (arg == "--list") { @@ -1555,6 +1814,9 @@ int main(int argc, char** argv) { if (correctness_only) { return run_linear_correctness(filter) ? 0 : 1; } + if (sdpa_correctness_only) { + return run_sdpa_correctness() ? 0 : 1; + } if ((linear || baseline) && !skip_correctness) { if (!run_linear_correctness(filter)) { std::cout << "correctness gate FAILED -- not running the perf sweep\n"; From 7de604dcb21f56a3df2516ca028efa1db431d57c Mon Sep 17 00:00:00 2001 From: Yanwen Xu Date: Fri, 21 Aug 2026 16:10:51 -0700 Subject: [PATCH 10/28] [ET-VK] Port shmem_double_buf4-tr's coopmat-staged A to 8da4w (UNVALIDATED) Adds linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr, a fork of the dbuf4 tile-sweep kernel that stages A through cooperative-matrix registers -- coopMatLoad from global, coopMatStore into LDS -- instead of a hand-rolled per-thread ivec4 copy. Ported from shaders/shmem_double_buf4-tr.comp on vk_cooperative_matrix_perf's gemm-ubm branch. Only the A half of the reference's idea is portable here. B is already column-major in LDS, so the reference's "tr" property is not a delta for it, and B cannot be coopmat-staged at all: the int4 weights need a nibble-extract that coopMatLoad cannot do, and a coopmat's per-lane layout is opaque so one cannot be built from unpacked registers. B staging is left byte-identical. A needed a new activation layout. The stock 4h4w packing is not row-major -- element [m4*K4 + k4] is an ivec4 whose component selects one of 4 rows, so the flat index is m4*(4*K4) + k4*4 + r, which is not affine in the row index and cannot be addressed by any RowMajor/ColumnMajor coopMatLoad (ColumnMajor also fails on contiguity, since a uint packs 4 K-values rather than 4 M-values). quantize_and_pack_4w_with_group_sums emits the same quantization and group sums into the existing kPackedInt8_4W layout instead, which is plain row-major. write_block in linear_int8_input_block.glslh assumes an ivec4-typed output, so it gets a SKIP_BLOCK_WRITE_HELPERS opt-out; no other shader is affected. Layout and kernel are chosen together. Both quantized_linear_impl (graph build time) and pick_linear_dqa_qw_shader (dispatch time) now go through dq8ca_coopmat_dispatch_eligible(), so a shape that falls back to the tiled path still gets the 4h4w layout that path expects. The predicate is guarded by dq8ca_variant_wants_rowmajor_a() first, so on the shipped default it short-circuits and nothing about the stock path changes. NOT the default and NOT validated on hardware: reachable only via ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4tr_txkgs. Verified so far only that it builds, that all 9 kernel variants and 16 packer variants land in spv.cpp, and that the SPIR-V matches the intent -- against dbuf4 at the same tile, OpCooperativeMatrixMulAddKHR is unchanged at 8 while LoadKHR goes 16 -> 24 and StoreKHR 4 -> 12, exactly A_TILES_PER_SG=4 times the prologue and main-loop staging sites. --- ...ar_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr.glsl | 699 ++++++++++++++++++ ...ar_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr.yaml | 141 ++++ .../ops/glsl/linear_int8_input_block.glslh | 9 + .../quantize_and_pack_4w_with_group_sums.glsl | 163 ++++ .../quantize_and_pack_4w_with_group_sums.yaml | 45 ++ .../graph/ops/impl/QuantizeDequantize.cpp | 75 +- .../graph/ops/impl/QuantizeDequantize.h | 14 + .../graph/ops/impl/QuantizedLinear.cpp | 169 +++-- 8 files changed, 1263 insertions(+), 52 deletions(-) create mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr.glsl create mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr.yaml create mode 100644 backends/vulkan/runtime/graph/ops/glsl/quantize_and_pack_4w_with_group_sums.glsl create mode 100644 backends/vulkan/runtime/graph/ops/glsl/quantize_and_pack_4w_with_group_sums.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr.glsl b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr.glsl new file mode 100644 index 00000000000..6027412c898 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr.glsl @@ -0,0 +1,699 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +/* + * "-tr" (coopmat-staged A) variant of linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4. + * + * Ported from shmem_double_buf4-tr.comp on vk_cooperative_matrix_perf's + * gemm-ubm branch. That reference file's delta over shmem_double_buf4.comp is + * that the global -> LDS staging goes through COOPERATIVE MATRIX REGISTERS + * (coopMatLoad from global -> coopmat<> array -> coopMatStore into shared) + * instead of a hand-rolled per-thread uvec4 copy, and that B lands in LDS + * column-major. + * + * Only the A half of that idea is portable to this kernel: + * + * - B is ALREADY column-major in LDS here (Bsh_int8 is K-contiguous per + * output column, read back with gl_CooperativeMatrixLayoutColumnMajor), + * so the reference's "tr" property is not a delta for B at all. + * - B CANNOT be coopmat-staged: the weights are int4, each ivec4 holding + * 8 columns x 4 K-values that need the nibble-extract / -8 / sign-pack + * below. coopMatLoad cannot unpack nibbles, and a coopmat's per-lane + * layout is opaque so one cannot be built from unpacked registers. + * B staging is therefore left byte-identical to dbuf4. + * - A CAN be coopmat-staged, but only against a ROW-MAJOR packed int8 + * activation buffer. The stock 4h4w layout (kPackedInt8_4H4W, produced by + * quantize_and_pack_4h4w_with_group_sums.glsl) is NOT row-major: element + * [m4 * K4 + k4] is an ivec4 whose COMPONENT selects one of 4 rows, so as + * a uint array the index is m4*(4*K4) + k4*4 + r, which is not affine in + * the row index and cannot be addressed by any RowMajor/ColumnMajor + * coopMatLoad. (ColumnMajor is out on contiguity too: a uint packs 4 + * K-values, not 4 M-values.) + * + * So this shader binds t_packed_int8_input as a SCALAR int array in the + * kPackedInt8_4W layout -- plain row-major int8, 4 K-values per int32, + * row stride K4 -- produced by quantize_and_pack_4w_with_group_sums.glsl. + * QuantizedLinear.cpp allocates that layout (and dispatches that packer) + * only when the active dq8ca variant is a "tsweep_dbuf4tr_t..." token AND + * the coopmat gate passes, so the tiled fallback never sees the wrong + * layout. Everything downstream of A staging -- LDS layout, int8 WMMA thread + * maps, group epilog, bias/store epilogue -- is unchanged from dbuf4. + * + * A staging (the actual -tr port): + * dbuf4: per-thread (m4, k4) ivec4 fetch; only A_ACTIVE_THREADS = + * (WG_TILE_M/4) * (WG_TILE_K/4) invocations participate, each + * scattering 4 rows into Ash_int8 with 4 scalar stores. + * dbuf4tr: per-SUBGROUP MMA_M x MMA_K tile fetch via coopMatLoad straight + * from global, then coopMatStore into the same Ash_int8 slot. + * The (WG_TILE_M/MMA_M) * (WG_TILE_K/MMA_K) tiles of a chunk are + * dealt round-robin across the NUM_SUBGROUPS subgroups. + * + * The loop structure is dbuf4's, unchanged: + * prologue: prefetch chunk 0 -> temp, store to slice 0 (no barrier) + * per iter: barrier -> prefetch(next) -> MMA(cur) -> store(next) + * and the nested `groups x chunks` loop with an unconditional group epilog is + * kept as-is (flattening it crashes the Xclipse PAL compiler at large + * spec-resolved trip counts -- see dbuf2's header). + * + * Selected at dispatch via + * ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4tr_txkgs<32|64> + * (QuantizedLinear.cpp), additive to the tsweep_dbuf4_t... and tsweep_t... + * namespaces. NOT the default -- unvalidated until it passes repeated + * test_llama_microbench --correctness-only runs (see dq8ca_coopmat_variant()'s + * comment on why a single pass is not proof). + * + * Performs: out[M,N] = dequant(int8_act) * dequant(int_w) (+ bias) + * via coopmat x coopmat -> coopmat on the matrix unit. + * + * Hard preconditions (in addition to dbuf4's): + * M % WG_TILE_M == 0, N % WG_TILE_N == 0, K % WG_TILE_K == 0, + * group_size % WG_TILE_K == 0, K % 4 == 0, + * WG_TILE_M % MMA_M == 0, WG_TILE_K % MMA_K == 0, + * t_packed_int8_input in kPackedInt8_4W (row-major) layout, + * device exposes coopmatx-> at 16x16x16. + */ + +#version 450 core + +#extension GL_KHR_cooperative_matrix : require +#extension GL_KHR_memory_scope_semantics : require +#extension GL_KHR_shader_subgroup_basic : enable +#extension GL_EXT_shader_explicit_arithmetic_types : require +#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_EXT_control_flow_attributes : enable + +#define PRECISION ${PRECISION} + +$if WEIGHT_NBITS == 4: + #define WEIGHT_INT4 + +$if HAS_BIAS: + #define HAS_BIAS + +$if WEIGHT_STORAGE == "buffer": + #define WEIGHT_BUFFER + +$if IO_STORAGE == "texture3d": + #define IO_TEXTURE + +layout(std430) buffer; + +#include "common.glslh" + +// Bindings — match add_linear_dqa_qw_node arg order: +// output(0), fp_input(1), packed_int8_input(2), int_input_sums(3 - unused), +// input_scales(4), input_zps(5), packed_weight(6), weight_sums(7), +// weight_scales(8), bias(9). +${layout_declare_tensor(B, "w", "t_output", "half", IO_STORAGE, is_scalar_array=True)} +// t_input is unread here -- the activations arrive already quantized in +// t_packed_int8_input -- but stays declared so the binding layout matches the +// dispatch site. It tracks IO_STORAGE so the two IO tensors stay consistent. +${layout_declare_tensor(B, "r", "t_input", "half", IO_STORAGE, is_scalar_array=False)} +// ROW-MAJOR (kPackedInt8_4W) packed activations: scalar int array, each +// element holding 4 K-contiguous int8, row stride K4 = K/4. This is the one +// binding that differs from dbuf4 (which takes the 4h4w ivec4 block layout); +// it is what makes the coopMatLoad-based A staging below addressable. +${layout_declare_tensor(B, "r", "t_packed_int8_input", "int", "buffer", is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_int8_input_sums", "int", "buffer", is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_int8_input_scales", "half", "texture3d")} +${layout_declare_tensor(B, "r", "t_int8_input_zps", "int8", "texture3d")} +${layout_declare_tensor(B, "r", "t_packed_weight", "int", WEIGHT_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_weight_sums", "int", "buffer", is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_weight_scales", "half", "buffer", is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_bias", "half", "buffer", is_scalar_array=True)} + +${layout_declare_ubo(B, "ivec4", "output_sizes")} +${layout_declare_ubo(B, "ivec4", "input_sizes")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +${layout_declare_spec_const(C, "int", "apply_bias", "0")} +// INT4 only; inert (0) for INT8 so the dispatcher's spec list lines up. +${layout_declare_spec_const(C, "int", "K4_per_group", "0")} +${layout_declare_spec_const(C, "int", "num_groups_arg", "0")} +${layout_declare_spec_const(C, "int", "out_N_arg", "0")} + +// Tile geometry +const uint MMA_M = ${MMA_M}; +const uint MMA_N = ${MMA_N}; +const uint MMA_K = ${MMA_K}; + +const uint WG_TILE_M = ${WG_TILE_M}; +const uint WG_TILE_N = ${WG_TILE_N}; +const uint WG_TILE_K = ${WG_TILE_K}; + +const uint SG_GRID_X = ${SG_GRID_X}; +const uint SG_GRID_Y = ${SG_GRID_Y}; +const uint SUBGROUP_SIZE = ${SUBGROUP_SIZE}; +const uint NUM_SUBGROUPS = SG_GRID_X * SG_GRID_Y; +const uint WG_SIZE = NUM_SUBGROUPS * SUBGROUP_SIZE; + +const uint SG_TILE_M = WG_TILE_M / SG_GRID_Y; +const uint SG_TILE_N = WG_TILE_N / SG_GRID_X; +const uint MMAS_PER_SG_M = SG_TILE_M / MMA_M; +const uint MMAS_PER_SG_N = SG_TILE_N / MMA_N; + +const uint A_SLAB_INT8 = WG_TILE_M * MMA_K; +const uint B_USEFUL_U32 = MMA_K / 4u; +const uint B_STRIDE_U32 = B_USEFUL_U32 + 1u; // +1 skew +const uint B_SLAB_U32 = WG_TILE_N * B_STRIDE_U32; +const uint NUM_K_SLABS = WG_TILE_K / MMA_K; + +const uint A_SLAB_U32 = A_SLAB_INT8 / 4u; +const uint A_STRIDE_U32 = MMA_K / 4u; + +// One ping-pong slice covers all K-slabs of one chunk. +const uint ASH_SLICE_U32 = NUM_K_SLABS * A_SLAB_U32; +const uint BSH_SLICE_U32 = NUM_K_SLABS * B_SLAB_U32; + +// Double-buffered MMA operand staging. +shared uint Ash_int8[2u * ASH_SLICE_U32]; +shared uint Bsh_int8[2u * BSH_SLICE_U32]; + +// Per-WG-tile-row activation params (loaded ONCE at WG start; constant +// across groups). +shared int izp_sh[WG_TILE_M]; // int32 (cast from int8 source) for broadcast +shared float ifs_sh[WG_TILE_M]; // float32 (cast from fp16 source) for broadcast + +// Per-(group, output-channel) weight params, ping-ponged by group parity. +// (For per-channel INT8 only slice 0 is ever used.) +shared int wsum_sh[2u * WG_TILE_N]; +shared float wsc_sh[2u * WG_TILE_N]; + +#ifdef HAS_BIAS +shared float bias_sh[WG_TILE_N]; +#endif + +#ifdef IO_TEXTURE +// Result staging for the imageStore epilogue, mirroring the fp16 kernel: +// SG_GRID_Y bands of MMA_M rows, each WG_TILE_N wide, row-major. A full +// WG_TILE_M x WG_TILE_N buffer would cost SG_GRID_Y/MMAS_PER_SG_M x more LDS +// and wreck occupancy. float16_t-typed because coopMatStore needs it. +const uint CSH_ROWS = SG_GRID_Y * MMA_M; +shared float16_t Csh_out[CSH_ROWS * WG_TILE_N]; +#endif + +// Running fp32 accumulator (across all groups). +coopmat + result[MMAS_PER_SG_M][MMAS_PER_SG_N]; + +// Per-group int32 MMA accumulator. +coopmat + accum_int32[MMAS_PER_SG_M][MMAS_PER_SG_N]; + +void main() { + const uvec2 tileID = uvec2(gl_WorkGroupID.xy); + const uvec2 warpInTile = uvec2( + gl_SubgroupID % SG_GRID_X, + gl_SubgroupID / SG_GRID_X); + + const uint K = uint(input_sizes.x); + const uint N = uint(output_sizes.x); + const uint N4 = (N + 3u) / 4u; + const uint nblocks_x_A = (K + 3u) >> 2u; + +#ifdef WEIGHT_INT4 + const uint num_groups = uint(num_groups_arg); + const uint CHUNKS_PER_GROUP = uint(K4_per_group) * 4u / WG_TILE_K; +#else + // Per-channel: a single quant "group" spanning all of K. The nested + // groups x chunks loop below collapses to a flat chunk loop, the wsum/wsc + // ping-pong never crosses a boundary, and the epilog runs exactly once. + const uint num_groups = 1u; + const uint CHUNKS_PER_GROUP = uint(num_groups_arg); +#endif + const uint num_chunks = num_groups * CHUNKS_PER_GROUP; + + const uint tile_m_start = WG_TILE_M * tileID.y; + const uint tile_n_start = WG_TILE_N * tileID.x; + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + result[i][j] = coopmat(0.0); + accum_int32[i][j] = coopmat(0); + } + } + + const uint K_BLOCKS_PER_CHUNK = WG_TILE_K >> 2u; + + // --- A staging tile map: one MMA_M x MMA_K coopmat tile per subgroup per + // slot. A chunk holds A_TILES_M x A_TILES_K such tiles; they are dealt + // round-robin across the NUM_SUBGROUPS subgroups, so every subgroup + // participates (dbuf4's per-thread map leaves WG_SIZE - + // A_ACTIVE_THREADS invocations idle whenever the tile is small). + // A_TILES_PER_SG rounds up, so the last slot may be partially used -- + // the `t < NUM_A_TILES` guard below is subgroup-uniform (t depends only + // on gl_SubgroupID), which is what coopmat ops require. + const uint A_TILES_M = WG_TILE_M / MMA_M; + const uint A_TILES_K = WG_TILE_K / MMA_K; // == NUM_K_SLABS + const uint NUM_A_TILES = A_TILES_M * A_TILES_K; + const uint A_TILES_PER_SG = (NUM_A_TILES + NUM_SUBGROUPS - 1u) / NUM_SUBGROUPS; + +#ifdef WEIGHT_INT4 + // --- B staging thread map: (block, col) slots; each slot extracts one + // ColumnMajor LDS uint (4 K-contiguous sign-extended int8) --- + const uint B_TOTAL_SLOTS = K_BLOCKS_PER_CHUNK * WG_TILE_N; + const uint B_SLOTS_PER_THREAD = B_TOTAL_SLOTS / WG_SIZE; + const uint N8_PER_TILE = WG_TILE_N >> 3u; +#else + // --- B staging thread map: one (k4, n4) ivec4 block per active thread --- + const uint B_FETCH_SLOTS = K_BLOCKS_PER_CHUNK * (WG_TILE_N >> 2u); + const uint N4_PER_TILE = WG_TILE_N >> 2u; + const uint b_k4_in_chunk = gl_LocalInvocationID.x / N4_PER_TILE; + const uint b_n_uint_col = gl_LocalInvocationID.x % N4_PER_TILE; + const bool b_active = gl_LocalInvocationID.x < B_FETCH_SLOTS; +#endif + + // Prefetch temp registers. temp_A is a coopmat array (the -tr change); + // indices into it are [[unroll]]-resolved compile-time constants, never + // dynamic -- dynamic indexing of a coopmat array is exactly the construct + // the Xclipse/AMD-PAL compiler has miscompiled before. + coopmat + temp_A[A_TILES_PER_SG]; +#ifdef WEIGHT_INT4 + ivec4 temp_B[B_SLOTS_PER_THREAD]; + int temp_wsum; + float temp_wsc; +#else + ivec4 temp_B; +#endif + + // ========================================================= + // PROLOGUE + // ========================================================= + if (gl_LocalInvocationID.x < (WG_TILE_M >> 2u)) { + const uint m4 = (tile_m_start >> 2u) + gl_LocalInvocationID.x; + const vec4 sc = vec4(texelFetch(t_int8_input_scales, ivec3(m4, 0, 0), 0)); + const ivec4 zp = texelFetch(t_int8_input_zps, ivec3(m4, 0, 0), 0); + const uint base = gl_LocalInvocationID.x * 4u; + ifs_sh[base + 0u] = sc.x; ifs_sh[base + 1u] = sc.y; + ifs_sh[base + 2u] = sc.z; ifs_sh[base + 3u] = sc.w; + izp_sh[base + 0u] = zp.x; izp_sh[base + 1u] = zp.y; + izp_sh[base + 2u] = zp.z; izp_sh[base + 3u] = zp.w; + } + // Group 0 weight sums/scales -> slice 0. + if (gl_LocalInvocationID.x < WG_TILE_N) { + const uint n_idx = tile_n_start + gl_LocalInvocationID.x; + f16vec4 sv = t_weight_scales[n_idx >> 2u]; + wsc_sh[gl_LocalInvocationID.x] = float(sv[n_idx & 3u]); + wsum_sh[gl_LocalInvocationID.x] = t_weight_sums[n_idx]; + } + memoryBarrierShared(); + barrier(); + + // izp/ifs are per-row activation params, constant across K groups — + // broadcast them into coopmats ONCE; the group epilog reuses them every + // group (they depend only on the row block i, not on the group or j). + coopmat + izp_bcast[MMAS_PER_SG_M]; + coopmat + ifs_bcast[MMAS_PER_SG_M]; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + const uint local_m_base = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + coopMatLoad( + izp_bcast[i], izp_sh, + local_m_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutColumnMajor); + coopMatLoad( + ifs_bcast[i], ifs_sh, + local_m_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutColumnMajor); + } + + // dbuf4: prefetch chunk 0 into temp registers, THEN store to slice 0 (no + // barrier here -- the main loop's first iteration barriers before + // reading slice 0). + [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { + const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; + if (t < NUM_A_TILES) { + const uint tm = t / A_TILES_K; + const uint tk = t % A_TILES_K; + // Offset and stride are in ARRAY ELEMENTS (int = 4 int8), matching how + // the MMA loop already addresses the shared uint arrays below. + coopMatLoad( + temp_A[s], t_packed_int8_input, + (tile_m_start + tm * MMA_M) * nblocks_x_A + tk * (MMA_K >> 2u), + nblocks_x_A, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#ifdef WEIGHT_INT4 + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { + const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; + const uint block_in_chunk = slot >> 3u; + const uint k4_blk = block_in_chunk / N8_PER_TILE; + const uint n8_blk = (tile_n_start >> 3u) + (block_in_chunk % N8_PER_TILE); +#ifdef WEIGHT_BUFFER + temp_B[si] = t_packed_weight[(n8_blk * nblocks_x_A) + k4_blk]; +#else + temp_B[si] = texelFetch(t_packed_weight, ivec2(k4_blk, n8_blk), 0); +#endif + } +#else + if (b_active) { + const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; +#ifdef WEIGHT_BUFFER + temp_B = t_packed_weight[(b_k4_in_chunk * N4) + block_x_w]; +#else + temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, b_k4_in_chunk), 0); +#endif + } +#endif + { + // store chunk 0 -> slice 0 + [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { + const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; + if (t < NUM_A_TILES) { + const uint tm = t / A_TILES_K; + const uint tk = t % A_TILES_K; + coopMatStore( + temp_A[s], Ash_int8, + tk * A_SLAB_U32 + (tm * MMA_M) * A_STRIDE_U32, + A_STRIDE_U32, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#ifdef WEIGHT_INT4 + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { + const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; + const uint block_in_chunk = slot >> 3u; + const uint col_in_block = slot & 7u; + const uint k4_in_chunk = block_in_chunk / N8_PER_TILE; + const uint n8_in_tile = block_in_chunk % N8_PER_TILE; + const uint r = col_in_block & 3u; + const uint parity = col_in_block >> 2u; + const int w = temp_B[si][r]; + const int base = int(4u * parity); + const int v0 = (((w >> (base + 0)) & 0xF) - 8) & 0xFF; + const int v1 = (((w >> (base + 8)) & 0xF) - 8) & 0xFF; + const int v2 = (((w >> (base + 16)) & 0xF) - 8) & 0xFF; + const int v3 = (((w >> (base + 24)) & 0xF) - 8) & 0xFF; + const uint n_col = n8_in_tile * 8u + r + parity * 4u; + const uint slab_idx = k4_in_chunk / (MMA_K >> 2u); + const uint k4_in_slab = k4_in_chunk % (MMA_K >> 2u); + Bsh_int8[slab_idx * B_SLAB_U32 + n_col * B_STRIDE_U32 + k4_in_slab] = + uint(v0 | (v1 << 8) | (v2 << 16) | (v3 << 24)); + } +#else + if (b_active) { + const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); + const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); + const uint n_col_base = b_n_uint_col * 4u; + [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { + Bsh_int8[slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab] = + uint(temp_B[n_in_blk]); + } + } +#endif + } + + // ========================================================= + // MAIN LOOP (dbuf4) — nested groups x chunks (kept nested; flattening it + // with a conditional coopmat epilog crashes the Xclipse PAL compiler at + // large spec-resolved trip counts). One barrier per chunk. Chunk + // iteration (global index `chunk`): + // 1. barrier — A/B slice (chunk%2) fully written; on the first chunk + // of group g, wsum/wsc slice (g%2) is too. + // 2. prefetch — chunk+1 (A blocks, B blocks) into temp; when chunk+1 + // starts a new group, also its wsum/wsc element. Skipped + // entirely on the final chunk. + // 3. int8 MMA — on slice (chunk%2) into accum_int32. + // 4. store — temp -> A/B slice ((chunk+1)%2), unpacking the weight; + // on a group boundary, wsum/wsc -> slice ((g+1)%2). + // The group epilog runs unconditionally at the tail of each group. + // ========================================================= + uint chunk = 0; + for (uint group_i = 0; group_i < num_groups; ++group_i) { + for (uint inner = 0; inner < CHUNKS_PER_GROUP; ++inner, ++chunk) { + const bool has_next = chunk + 1u < num_chunks; + const bool group_crossing = has_next && (inner + 1u == CHUNKS_PER_GROUP); + const uint cur_a = (chunk % 2u) * ASH_SLICE_U32; + const uint cur_b = (chunk % 2u) * BSH_SLICE_U32; + const uint nxt_a = ((chunk + 1u) % 2u) * ASH_SLICE_U32; + const uint nxt_b = ((chunk + 1u) % 2u) * BSH_SLICE_U32; + + barrier(); + + // --- 2. prefetch chunk+1 -> temp --- + if (has_next) { + const uint chunkK_nxt = (chunk + 1u) * WG_TILE_K; + [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { + const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; + if (t < NUM_A_TILES) { + const uint tm = t / A_TILES_K; + const uint tk = t % A_TILES_K; + coopMatLoad( + temp_A[s], t_packed_int8_input, + (tile_m_start + tm * MMA_M) * nblocks_x_A + + (chunkK_nxt >> 2u) + tk * (MMA_K >> 2u), + nblocks_x_A, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#ifdef WEIGHT_INT4 + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { + const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; + const uint block_in_chunk = slot >> 3u; + const uint k4_blk = (chunkK_nxt >> 2u) + block_in_chunk / N8_PER_TILE; + const uint n8_blk = (tile_n_start >> 3u) + (block_in_chunk % N8_PER_TILE); +#ifdef WEIGHT_BUFFER + temp_B[si] = t_packed_weight[(n8_blk * nblocks_x_A) + k4_blk]; +#else + temp_B[si] = texelFetch(t_packed_weight, ivec2(k4_blk, n8_blk), 0); +#endif + } + if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { + const uint n_idx = tile_n_start + gl_LocalInvocationID.x; + f16vec4 sv = t_weight_scales[(group_i + 1u) * N4 + (n_idx >> 2u)]; + temp_wsc = float(sv[n_idx & 3u]); + temp_wsum = t_weight_sums[(group_i + 1u) * N + n_idx]; + } +#else + if (b_active) { + const uint block_y_w = (chunkK_nxt >> 2u) + b_k4_in_chunk; + const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; +#ifdef WEIGHT_BUFFER + temp_B = t_packed_weight[(block_y_w * N4) + block_x_w]; +#else + temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, block_y_w), 0); +#endif + } +#endif + } + + // --- 3. int8 MMA on the cur slice --- + [[unroll]] for (uint k = 0; k < NUM_K_SLABS; ++k) { + const uint slab_a_base_u32 = cur_a + k * A_SLAB_U32; + const uint slab_b_base_u32 = cur_b + k * B_SLAB_U32; + + coopmat matA[MMAS_PER_SG_M]; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + const uint row_a = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + coopMatLoad( + matA[i], Ash_int8, + slab_a_base_u32 + row_a * A_STRIDE_U32, + A_STRIDE_U32, + gl_CooperativeMatrixLayoutRowMajor); + } + + coopmat matB; + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint col_b = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopMatLoad( + matB, Bsh_int8, + slab_b_base_u32 + col_b * B_STRIDE_U32, + B_STRIDE_U32, + gl_CooperativeMatrixLayoutColumnMajor); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + accum_int32[i][j] = coopMatMulAdd(matA[i], matB, accum_int32[i][j]); + } + } + } + + // --- 4. store temp (chunk+1) -> nxt slice --- + if (has_next) { + [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { + const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; + if (t < NUM_A_TILES) { + const uint tm = t / A_TILES_K; + const uint tk = t % A_TILES_K; + coopMatStore( + temp_A[s], Ash_int8, + nxt_a + tk * A_SLAB_U32 + (tm * MMA_M) * A_STRIDE_U32, + A_STRIDE_U32, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#ifdef WEIGHT_INT4 + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { + const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; + const uint block_in_chunk = slot >> 3u; + const uint col_in_block = slot & 7u; + const uint k4_in_chunk = block_in_chunk / N8_PER_TILE; + const uint n8_in_tile = block_in_chunk % N8_PER_TILE; + const uint r = col_in_block & 3u; + const uint parity = col_in_block >> 2u; + const int w = temp_B[si][r]; + const int base = int(4u * parity); + const int v0 = (((w >> (base + 0)) & 0xF) - 8) & 0xFF; + const int v1 = (((w >> (base + 8)) & 0xF) - 8) & 0xFF; + const int v2 = (((w >> (base + 16)) & 0xF) - 8) & 0xFF; + const int v3 = (((w >> (base + 24)) & 0xF) - 8) & 0xFF; + const uint n_col = n8_in_tile * 8u + r + parity * 4u; + const uint slab_idx = k4_in_chunk / (MMA_K >> 2u); + const uint k4_in_slab = k4_in_chunk % (MMA_K >> 2u); + Bsh_int8[nxt_b + slab_idx * B_SLAB_U32 + n_col * B_STRIDE_U32 + k4_in_slab] = + uint(v0 | (v1 << 8) | (v2 << 16) | (v3 << 24)); + } + if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { + const uint wbase_nxt = ((group_i + 1u) % 2u) * WG_TILE_N; + wsum_sh[wbase_nxt + gl_LocalInvocationID.x] = temp_wsum; + wsc_sh[wbase_nxt + gl_LocalInvocationID.x] = temp_wsc; + } +#else + if (b_active) { + const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); + const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); + const uint n_col_base = b_n_uint_col * 4u; + [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { + Bsh_int8[nxt_b + slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab] = + uint(temp_B[n_in_blk]); + } + } +#endif + } + } // chunks + + // --- Group epilog: dequant accum_int32 -> result, reset accum --- + { + const uint wbase = (group_i % 2u) * WG_TILE_N; + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + + coopmat wsum_bcast; + coopMatLoad( + wsum_bcast, wsum_sh, + wbase + local_n_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutRowMajor); + + coopmat wsc_bcast; + coopMatLoad( + wsc_bcast, wsc_sh, + wbase + local_n_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutRowMajor); + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + coopmat adjusted = + accum_int32[i][j] - izp_bcast[i] * wsum_bcast; + coopmat adjusted_fp = + coopmat(adjusted); + coopmat scales_outer = + ifs_bcast[i] * wsc_bcast; + result[i][j] += adjusted_fp * scales_outer; + accum_int32[i][j] = coopmat(0); + } + } + } + } // groups + + // --- Bias (optional) --- +#ifdef HAS_BIAS + if (apply_bias > 0) { + for (uint t = gl_LocalInvocationID.x; t < WG_TILE_N; t += WG_SIZE) { + bias_sh[t] = float(t_bias[tile_n_start + t]); + } + memoryBarrierShared(); + barrier(); + } +#endif + + // --- Store result tile --- + // N for the store address math MUST come from the spec constant, not the + // sizes UBO (see out_N_arg above). +#ifdef IO_TEXTURE + // Epilogue iteration i drains accumulator row-block i from EVERY subgroup + // into Csh_out at once, so the SG_GRID_Y bands it holds are disjoint global + // row ranges; the whole workgroup then imageStores them. lr / MMA_M is the + // writing subgroup's warpInTile.y, so the global row reproduces the buffer + // path's gi exactly. + // + // PORTABILITY NOTE: the barrier() in the loop body keeps this loop rolled + // despite [[unroll]], so result[i][j] IS dynamically indexed. Coopmat arrays + // are opaque per-lane storage and dynamic indexing is exactly the construct + // the Xclipse/AMD-PAL compiler has broken before -- check this first if the + // texture variants miscompile on M51. + const uint CSH_TEXELS_PER_ROW = WG_TILE_N / 4u; + const uint CSH_TEXELS = CSH_ROWS * CSH_TEXELS_PER_ROW; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + // Guards Csh_out against the previous iteration's readers. Inert on i == 0 + // but must stay unconditional to remain workgroup-uniform. + barrier(); + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { +#ifdef HAS_BIAS + if (apply_bias > 0) { + const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopmat bias_tile; + coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); + result[i][j] += bias_tile; + } +#endif + coopmat out_tile = + coopmat(result[i][j]); + coopMatStore( + out_tile, Csh_out, + warpInTile.y * MMA_M * WG_TILE_N + + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j), + WG_TILE_N, + gl_CooperativeMatrixLayoutRowMajor); + } + memoryBarrierShared(); + barrier(); + + for (uint t = gl_LocalInvocationID.x; t < CSH_TEXELS; t += WG_SIZE) { + const uint lr = t / CSH_TEXELS_PER_ROW; + const uint lc4 = t % CSH_TEXELS_PER_ROW; + const uint m = + tile_m_start + (lr / MMA_M) * SG_TILE_M + i * MMA_M + (lr % MMA_M); + const uint base = lr * WG_TILE_N + lc4 * 4u; + imageStore( + t_output, + ivec3(tile_n_start / 4u + lc4, m, 0), + vec4( + float(Csh_out[base]), + float(Csh_out[base + 1u]), + float(Csh_out[base + 2u]), + float(Csh_out[base + 3u]))); + } + } +#else + const uint N_out = uint(out_N_arg); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint gi = tile_m_start + MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + const uint gj = tile_n_start + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + +#ifdef HAS_BIAS + if (apply_bias > 0) { + const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopmat bias_tile; + coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); + result[i][j] += bias_tile; + } +#endif + + coopmat out_tile = + coopmat(result[i][j]); + coopMatStore( + out_tile, t_output, + gi * N_out + gj, N_out, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#endif // IO_TEXTURE +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr.yaml b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr.yaml new file mode 100644 index 00000000000..e7ee2bf9d9e --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr.yaml @@ -0,0 +1,141 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# "-tr" (coopmat-staged A) variants of the int8 dq8ca_q4gsw coopmat kernel, +# ported from shmem_double_buf4-tr.comp (vk_cooperative_matrix_perf, +# gemm-ubm). linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr.glsl is a fork of +# linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4.glsl with ONLY the A-side global -> +# LDS staging swapped from a per-thread ivec4 copy to +# coopMatLoad(global) -> coopmat<> -> coopMatStore(LDS). B staging is +# byte-identical (int4 nibble unpack cannot be coopmat-staged) and B was +# already column-major in LDS, so it is not a delta. +# +# These variants require t_packed_int8_input in the ROW-MAJOR +# kPackedInt8_4W layout (produced by quantize_and_pack_4w_with_group_sums), +# NOT the stock 4h4w block layout. QuantizedLinear.cpp switches the tensor +# layout and the packer node together with the variant token, and only when +# the coopmat gate passes -- see dq8ca_wants_rowmajor_int8_input(). +# +# Selected via +# ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4tr_txkgs. +# NOT the shipped default: unvalidated until it passes repeated +# test_llama_microbench --correctness-only runs (a single pass is not proof -- +# see dq8ca_coopmat_variant()'s comment). +# +# Tile-geometry preconditions beyond dbuf4's: WG_TILE_M % MMA_M == 0 and +# WG_TILE_K % MMA_K == 0 (both hold for every entry below). The seed set is +# deliberately small -- widen it from a sweep once correctness is established. + +linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr: + parameter_names_with_default_values: + PRECISION: highp + HAS_BIAS: false + IO_STORAGE: buffer + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + MMA_M: 16 + MMA_N: 16 + MMA_K: 16 + WG_TILE_M: 64 + WG_TILE_N: 32 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + shader_variants: + # production 8da4w tile (dbuf2/dbuf4 default) -- the A/B anchor + # A tiles/chunk = 8, subgroups = 2 -> A_TILES_PER_SG = 4 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr_t64x32k32g12s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 32 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr_t64x32k32g12s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 32 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr_t64x32k32g12s64_buffer_buffer_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: buffer + WG_TILE_M: 64 + WG_TILE_N: 32 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + # q4gsw production geometry; 8 A-tiles over 4 subgroups + # A tiles/chunk = 8, subgroups = 4 -> A_TILES_PER_SG = 2 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr_t128x128k16g22s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr_t128x128k16g22s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + # 1 A-tile per subgroup (A_TILES_PER_SG == 1, no leftover slot) + # A tiles/chunk = 8, subgroups = 8 -> A_TILES_PER_SG = 1 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr_t64x64k32g24s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr_t64x64k32g24s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + # 16 A-tiles over 8 subgroups + # A tiles/chunk = 16, subgroups = 8 -> A_TILES_PER_SG = 2 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr_t128x64k32g24s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr_t128x64k32g24s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_int8_input_block.glslh b/backends/vulkan/runtime/graph/ops/glsl/linear_int8_input_block.glslh index 8f19418cd19..e2b58514480 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_int8_input_block.glslh +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_int8_input_block.glslh @@ -69,6 +69,14 @@ void quantize_and_pack( } } +// The 4h4w block writer. Guarded so that a shader which reuses the quantize +// helpers above but emits a DIFFERENT output layout can opt out: it assumes +// t_packed_int8_input is an ivec4-typed (block-packed) resource, which would +// fail to compile against a scalar int array even if never called. Define +// SKIP_BLOCK_WRITE_HELPERS before including this file to suppress it -- see +// quantize_and_pack_4w_with_group_sums.glsl, which writes the row-major +// kPackedInt8_4W layout instead. +#ifndef SKIP_BLOCK_WRITE_HELPERS void write_block( const Int8InputBlock block, const int block_x, @@ -80,5 +88,6 @@ void write_block( imageStore(t_packed_int8_input, ivec3(block_x, block_y, 0), block.data); #endif // OUTPUT_BUFFER } +#endif // SKIP_BLOCK_WRITE_HELPERS #endif // LINEAR_INT8_INPUT_BLOCK_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/glsl/quantize_and_pack_4w_with_group_sums.glsl b/backends/vulkan/runtime/graph/ops/glsl/quantize_and_pack_4w_with_group_sums.glsl new file mode 100644 index 00000000000..85bd48416c2 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/quantize_and_pack_4w_with_group_sums.glsl @@ -0,0 +1,163 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +// ROW-MAJOR (kPackedInt8_4W) counterpart of +// quantize_and_pack_4h4w_with_group_sums.glsl. +// +// Identical quantization and per-group sum reduction; the ONLY difference is +// the output layout. The 4h4w version writes an ivec4 per (m4, k4) block whose +// COMPONENT selects one of 4 rows -- as a scalar array that is +// index = m4*(4*K4) + k4*4 + r, which is not affine in the row index and so +// cannot be addressed by a cooperative-matrix load. This version writes each +// of the 4 quantized rows to its own row-major slot, +// t_packed_int8_input[m * K4 + k4], giving plain row-major int8 with 4 +// K-contiguous values per int32 and a row stride of K4. +// +// That is exactly what linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr.glsl needs to +// stage A through coopMatLoad/coopMatStore. Nothing else consumes this layout; +// QuantizedLinear.cpp dispatches this packer (and allocates the tensor as +// kPackedInt8_4W) only when a tsweep_dbuf4tr_t... variant is active AND the +// coopmat gate passes, so the tiled fallback never sees it. +// +// Buffer output only -- the coopmat A staging reads an SSBO. + +#version 450 core + +${define_required_extensions(INPUT_STORAGE, DTYPE)} +${define_required_extensions("texture3d", "int8")} + +#define PRECISION ${PRECISION} +#define VEC4_T ${texel_load_type(DTYPE, INPUT_STORAGE)} +#define T ${texel_load_component_type(DTYPE, INPUT_STORAGE)} + +$if OUTPUT_STORAGE == "buffer": + #define OUTPUT_BUFFER +$if INPUT_STORAGE == "buffer": + #define INPUT_BUFFER + +#extension GL_EXT_integer_dot_product : require + +#define NUM_GROUPS_PER_WG ${NUM_GROUPS_PER_WG} +#define NUM_WORKERS_PER_GROUP ${NUM_WORKERS_PER_GROUP} + +layout(std430) buffer; + +#include "common.glslh" + +${layout_declare_tensor(B, "w", "t_packed_int8_input", "int", OUTPUT_STORAGE, is_scalar_array=True)} +${layout_declare_tensor(B, "w", "t_int8_input_sums", "int", "buffer", is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_input", DTYPE, INPUT_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_int8_input_scales", DTYPE, "texture3d")} +${layout_declare_tensor(B, "r", "t_int8_input_zps", "int8" if ZP_DTYPE_MODE == "zpint8" else DTYPE, "texture3d")} + +${layout_declare_ubo(B, "ivec4", "input_sizes")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +${layout_declare_spec_const(C, "int", "K4_per_group", "0")} + +shared ivec4 shared_sums[NUM_GROUPS_PER_WG][NUM_WORKERS_PER_GROUP]; + +#define TILE_M4 1 +#define TILE_K4 1 + +#define TILE_M 4 + +// The shared helper's write_block() assumes an ivec4-typed output resource; +// this shader's output is a scalar int array, so opt out and write our own. +#define SKIP_BLOCK_WRITE_HELPERS +#include "linear_int8_input_block.glslh" + +// Scatters the 4 rows of a quantized block into row-major slots. Unlike the +// 4h4w layout, kPackedInt8_4W has NO row padding (outer_dim_align == 1), so +// the tail rows of a workgroup whose m4 block runs past M must be dropped -- +// hence the explicit bound check that write_block() does not need. +void write_block_rowmajor( + const Int8InputBlock block, + const int k4, + const int m, + const int K4, + const int M) { + for (int row = 0; row < 4; ++row) { + const int m_row = m + row; + if (m_row < M) { + t_packed_int8_input[m_row * K4 + k4] = block.data[row]; + } + } +} +#include "linear_int8_input_scales_zps_load.glslh" +#include "linear_fp_input_tile_load.glslh" + +void main() { + const int group_idx = int(gl_GlobalInvocationID.x); + const int m4 = int(gl_GlobalInvocationID.y); + + const int worker_id = int(gl_LocalInvocationID.z); + const int group_offset = int(gl_LocalInvocationID.x); + + const int K = input_sizes.x; + const int M = input_sizes.y; + + // K4 and M4 represent the number of blocks in each dimension. + const int K4 = div_up_4(K); + const int M4 = div_up_4(M); + + const int num_groups = K4 / K4_per_group;; + + if (group_idx >= num_groups || m4 >= M4) { + return; + } + + const int start_k4 = group_idx * K4_per_group + worker_id; + const int end_k4 = (group_idx + 1) * K4_per_group; + + Int8InputScales input_scales; + Int8InputZeroPoints input_zps; + load_int8_input_scales_and_zps(input_scales, input_zps, m4); + + // row of the input tensor to start loading from + const int m = mul_4(m4); + + FPInputTile in_tile; + Int8InputBlock packed_block; + + ivec4 local_sum = ivec4(0, 0, 0, 0); + const int packed_ones = 0x01010101; + + for (int k4 = start_k4; k4 < end_k4; k4 += NUM_WORKERS_PER_GROUP) { + load_input_tile_no_checks(in_tile, k4, m, K4, M); + quantize_and_pack(packed_block, in_tile, input_scales, input_zps); + + // Sum the quantized values in the block + [[unroll]] for (int m = 0; m < TILE_M; m++) { + local_sum[m] += dotPacked4x8AccSatEXT( + packed_block.data[m], packed_ones, local_sum[m]); + } + write_block_rowmajor(packed_block, k4, m, K4, M); + } + + shared_sums[group_offset][worker_id] = local_sum; + + memoryBarrierShared(); + barrier(); + + // Tree reduction to compute the overall result + for (int i = NUM_WORKERS_PER_GROUP / 2; i > 0; i >>= 1) { + if (worker_id < i) { + shared_sums[group_offset][worker_id] = + shared_sums[group_offset][worker_id] + + shared_sums[group_offset][worker_id + i]; + } + memoryBarrierShared(); + barrier(); + } + + if (worker_id == 0) { + t_int8_input_sums[group_idx * M4 + m4] = shared_sums[group_offset][0]; + } +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/quantize_and_pack_4w_with_group_sums.yaml b/backends/vulkan/runtime/graph/ops/glsl/quantize_and_pack_4w_with_group_sums.yaml new file mode 100644 index 00000000000..533a0fb4ee5 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/quantize_and_pack_4w_with_group_sums.yaml @@ -0,0 +1,45 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# Row-major (kPackedInt8_4W) counterpart of +# quantize_and_pack_4h4w_with_group_sums. Same quantization, same per-group +# sum reduction, same workgroup shapes -- only the output layout differs. +# Consumed exclusively by linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr, whose +# coopmat-staged A path needs a row-major int8 activation buffer that a +# coopMatLoad can address. Buffer output only. +# +# The variant matrix mirrors the 4h4w yaml minus the buffer-output-only +# restriction, so pick_quantize_and_pack_4w_with_group_sums_shader() can build +# the same name suffixes (_o2w32/_o4w16 + storage + dtype + zp mode). + +quantize_and_pack_4w_with_group_sums: + parameter_names_with_default_values: + DTYPE: float + OUTPUT_STORAGE: buffer + INPUT_STORAGE: texture3d + NUM_GROUPS_PER_WG: 2 + NUM_WORKERS_PER_GROUP: 32 + ZP_DTYPE_MODE: zpint8 + generate_variant_forall: + DTYPE: + - VALUE: half + - VALUE: float + ZP_DTYPE_MODE: + - VALUE: zpint8 + - VALUE: zpinherit + shader_variants: + - NAME: quantize_and_pack_4w_with_group_sums_o2w32_buffer_texture3d + - NAME: quantize_and_pack_4w_with_group_sums_o2w32_buffer_buffer + OUTPUT_STORAGE: buffer + INPUT_STORAGE: buffer + - NAME: quantize_and_pack_4w_with_group_sums_o4w16_buffer_texture3d + NUM_GROUPS_PER_WG: 4 + NUM_WORKERS_PER_GROUP: 16 + - NAME: quantize_and_pack_4w_with_group_sums_o4w16_buffer_buffer + NUM_GROUPS_PER_WG: 4 + NUM_WORKERS_PER_GROUP: 16 + OUTPUT_STORAGE: buffer + INPUT_STORAGE: buffer diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizeDequantize.cpp b/backends/vulkan/runtime/graph/ops/impl/QuantizeDequantize.cpp index 98f97eab572..dc545cca0c8 100644 --- a/backends/vulkan/runtime/graph/ops/impl/QuantizeDequantize.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/QuantizeDequantize.cpp @@ -60,10 +60,11 @@ utils::uvec3 quantize_and_pack_4h4w_global_wg_size( 1u}; } -vkapi::ShaderInfo pick_quantize_and_pack_4h4w_with_group_sums_shader( +static vkapi::ShaderInfo pick_quantize_and_pack_with_group_sums_shader( ComputeGraph* graph, const std::vector& args, - const std::vector& resize_args) { + const std::vector& resize_args, + const char* base_name) { const ValueRef packed_int_input = args.at(0).refs.at(0); const ValueRef fp_input = args.at(1).refs.at(0); const ValueRef packed_input_zps = args.at(1).refs.at(2); @@ -71,7 +72,7 @@ vkapi::ShaderInfo pick_quantize_and_pack_4h4w_with_group_sums_shader( const int64_t group_size_val = graph->extract_scalar(group_size); - std::string shader_name = "quantize_and_pack_4h4w_with_group_sums"; + std::string shader_name = base_name; if (group_size_val >= 128) { shader_name += "_o2w32"; } else { @@ -87,6 +88,25 @@ vkapi::ShaderInfo pick_quantize_and_pack_4h4w_with_group_sums_shader( return VK_KERNEL_FROM_STR(shader_name); } +vkapi::ShaderInfo pick_quantize_and_pack_4h4w_with_group_sums_shader( + ComputeGraph* graph, + const std::vector& args, + const std::vector& resize_args) { + return pick_quantize_and_pack_with_group_sums_shader( + graph, args, resize_args, "quantize_and_pack_4h4w_with_group_sums"); +} + +// Row-major (kPackedInt8_4W) output layout; same quantization and group-sum +// reduction, same workgroup shapes, so the two share every dispatch helper +// below. Only linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr consumes it. +vkapi::ShaderInfo pick_quantize_and_pack_4w_with_group_sums_shader( + ComputeGraph* graph, + const std::vector& args, + const std::vector& resize_args) { + return pick_quantize_and_pack_with_group_sums_shader( + graph, args, resize_args, "quantize_and_pack_4w_with_group_sums"); +} + utils::uvec3 pick_quantize_and_pack_4h4w_with_group_sums_global_wg_size( ComputeGraph* graph, const vkapi::ShaderInfo& shader, @@ -198,7 +218,7 @@ void add_quantize_and_pack_4h4w_node( {})); } -void add_quantize_and_pack_4h4w_with_group_sums_node( +static void add_quantize_and_pack_with_group_sums_node( ComputeGraph& graph, const QuantizationConfig& input_quant_config, const ValueRef fp_input, @@ -206,7 +226,8 @@ void add_quantize_and_pack_4h4w_with_group_sums_node( const ValueRef packed_input_scales, const ValueRef packed_input_zps, const ValueRef packed_int_input, - const ValueRef group_size) { + const ValueRef group_size, + const DynamicDispatchNode::PickShaderFn& pick_shader) { // Only certain quantization types supported at the moment VK_CHECK_COND(input_quant_config.granularity == kPerChannel); @@ -221,7 +242,7 @@ void add_quantize_and_pack_4h4w_with_group_sums_node( graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, - pick_quantize_and_pack_4h4w_with_group_sums_shader, + pick_shader, pick_quantize_and_pack_4h4w_with_group_sums_global_wg_size, pick_quantize_and_pack_4h4w_with_group_sums_local_wg_size, // Inputs and Outputs @@ -237,6 +258,48 @@ void add_quantize_and_pack_4h4w_with_group_sums_node( {group_size})); } +void add_quantize_and_pack_4h4w_with_group_sums_node( + ComputeGraph& graph, + const QuantizationConfig& input_quant_config, + const ValueRef fp_input, + const ValueRef int_input_sums, + const ValueRef packed_input_scales, + const ValueRef packed_input_zps, + const ValueRef packed_int_input, + const ValueRef group_size) { + add_quantize_and_pack_with_group_sums_node( + graph, + input_quant_config, + fp_input, + int_input_sums, + packed_input_scales, + packed_input_zps, + packed_int_input, + group_size, + pick_quantize_and_pack_4h4w_with_group_sums_shader); +} + +void add_quantize_and_pack_4w_with_group_sums_node( + ComputeGraph& graph, + const QuantizationConfig& input_quant_config, + const ValueRef fp_input, + const ValueRef int_input_sums, + const ValueRef packed_input_scales, + const ValueRef packed_input_zps, + const ValueRef packed_int_input, + const ValueRef group_size) { + add_quantize_and_pack_with_group_sums_node( + graph, + input_quant_config, + fp_input, + int_input_sums, + packed_input_scales, + packed_input_zps, + packed_int_input, + group_size, + pick_quantize_and_pack_4w_with_group_sums_shader); +} + // // Dispatch utilities (Conv2d) // diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizeDequantize.h b/backends/vulkan/runtime/graph/ops/impl/QuantizeDequantize.h index 96e9cc7c1d3..0d9d4728952 100644 --- a/backends/vulkan/runtime/graph/ops/impl/QuantizeDequantize.h +++ b/backends/vulkan/runtime/graph/ops/impl/QuantizeDequantize.h @@ -44,6 +44,20 @@ void add_quantize_and_pack_4h4w_with_group_sums_node( const ValueRef packed_int_input, const ValueRef group_size); +// Same as above but emits the row-major kPackedInt8_4W layout instead of the +// 4h4w block layout. Only the coopmat-staged-A dq8ca kernel +// (linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr) can consume it; packed_int_input +// must have been allocated with utils::kPackedInt8_4W. +void add_quantize_and_pack_4w_with_group_sums_node( + ComputeGraph& graph, + const QuantizationConfig& input_quant_config, + const ValueRef fp_input, + const ValueRef int_input_sums, + const ValueRef packed_input_scales, + const ValueRef packed_input_zps, + const ValueRef packed_int_input, + const ValueRef group_size); + // // Quantize, Dequantize for Convolution // diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp index 6e059ae8f5c..ab9201d8af8 100644 --- a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp @@ -94,6 +94,12 @@ static const char* const kTsweepPrefixes[] = { "tsweep_dbuf2_t", "tsweep_dbuf3_t", "tsweep_dbuf4_t", + // "-tr": dbuf4's loop with the A-side global -> LDS staging swapped to + // coopMatLoad/coopMatStore (ported from shmem_double_buf4-tr.comp). Stays + // mutually exclusive with "tsweep_dbuf4_t" because position 12 is 't' vs + // '_', so prefix order here does not matter. dq8ca only -- there is no + // q4gsw "-tr" shader. + "tsweep_dbuf4tr_t", "tsweep_t", }; @@ -426,6 +432,59 @@ static bool can_use_q4gsw_coopmat( return true; } +// True when the active dq8ca variant is a "-tr" token, i.e. the coopmat-staged +// A kernel. That kernel binds t_packed_int8_input as a scalar int array in the +// row-major kPackedInt8_4W layout instead of the 4h4w ivec4 block layout, +// because no coopMatLoad can address 4h4w (its component index selects a row, +// making the flat index non-affine in the row). +static bool dq8ca_variant_wants_rowmajor_a() { + return dq8ca_coopmat_variant().rfind("tsweep_dbuf4tr_t", 0) == 0; +} + +// Mirrors the coopmat branch of pick_linear_dqa_qw_shader() so graph-build time +// (which picks the activation layout and the packer node) and dispatch time +// (which picks the kernel) can never disagree. Both call this: if it returns +// true the dq8ca coopmat kernel WILL be selected, so it is safe to hand it the +// row-major layout; if it returns false the tiled fallback runs and must get +// the stock 4h4w layout. Both call sites read the same immutable graph state +// and the same process-level variant token, so they cannot drift. +static bool dq8ca_coopmat_dispatch_eligible( + ComputeGraph* graph, + const ValueRef output, + const ValueRef fp_input, + const ValueRef bias_data, + const int64_t group_size) { + if (is_gemv(graph, fp_input)) { + return false; + } + if (!graph->context()->adapter_ptr()->supports_int8_cooperative_matrix()) { + return false; + } + // The alignment gate must use the ACTIVE sweep variant's own tile dims, not + // the shipped default's (same rationale as the q4gsw tsweep hook). + // + // allow_texture_io: the dq8ca texture-IO shader declares t_input (fp_input) + // with an IO_STORAGE-typed binding even though it is never read in the + // shader body (activations arrive pre-quantized in t_packed_int8_input + // instead) -- Vulkan still requires the bound resource's storage type to + // match the declared binding type, so fp_input must genuinely be texture3d + // too when texture IO is active. Same requirement as q4gsw; no separate + // check needed. + const CoopmatTileDims dims = + parse_tsweep_tile(dq8ca_coopmat_variant(), kDq8caQ4gswCoopmatDims); + return can_use_q4gsw_coopmat( + graph, + output, + fp_input, + group_size, + bias_data, + dims.m, + dims.n, + dims.k, + /*allow_texture_io=*/true, + dims.sg_grid_y); +} + vkapi::ShaderInfo pick_linear_qw_shader( ComputeGraph* graph, const std::vector& args, @@ -513,43 +572,25 @@ vkapi::ShaderInfo pick_linear_dqa_qw_shader( // Use the coopmat shader for 4-bit dq8ca dispatches when the device // enumerates VK_COMPONENT_TYPE_SINT8_KHR in its cooperative matrix property - // list and the shape aligns; tiled otherwise. - if (weight_is_4bit && !is_gemv_case && - graph->context()->adapter_ptr()->supports_int8_cooperative_matrix()) { - const int64_t group_size = - graph->extract_scalar(resize_args.at(0)); - // Alignment gate must use the ACTIVE sweep variant's own tile dims (same - // rationale as the q4gsw tsweep hook above). - const CoopmatTileDims active_dq8ca_dims = - parse_tsweep_tile(dq8ca_coopmat_variant(), kDq8caQ4gswCoopmatDims); - // The dq8ca texture-IO shader declares t_input (fp_input) with an - // IO_STORAGE-typed binding even though it's never read in the shader body - // (activations arrive pre-quantized in t_packed_int8_input instead) -- - // Vulkan still requires the bound resource's storage type to match the - // declared binding type, so fp_input must genuinely be texture3d too when - // texture IO is active. Same requirement as q4gsw; no separate check - // needed. - if (can_use_q4gsw_coopmat( - graph, - out, - fp_input, - group_size, - resize_args.at(2), - active_dq8ca_dims.m, - active_dq8ca_dims.n, - active_dq8ca_dims.k, - /*allow_texture_io=*/true, - active_dq8ca_dims.sg_grid_y)) { - std::string kernel_name = "linear_dq8ca_q4gsw_coopmat"; - const std::string& dq8ca_variant = dq8ca_coopmat_variant(); - if (!dq8ca_variant.empty()) { - kernel_name += "_" + dq8ca_variant; - } - add_storage_type_suffix(kernel_name, graph->storage_type_of(out)); - add_storage_type_suffix(kernel_name, graph->storage_type_of(int_weight)); - add_dtype_suffix(kernel_name, graph->dtype_of(out)); - return VK_KERNEL_FROM_STR(kernel_name); + // list and the shape aligns; tiled otherwise. The eligibility test lives in + // dq8ca_coopmat_dispatch_eligible() because quantized_linear_impl() has to + // ask the same question at graph-build time to pick the activation layout. + if (weight_is_4bit && + dq8ca_coopmat_dispatch_eligible( + graph, + out, + fp_input, + resize_args.at(2), + graph->extract_scalar(resize_args.at(0)))) { + std::string kernel_name = "linear_dq8ca_q4gsw_coopmat"; + const std::string& dq8ca_variant = dq8ca_coopmat_variant(); + if (!dq8ca_variant.empty()) { + kernel_name += "_" + dq8ca_variant; } + add_storage_type_suffix(kernel_name, graph->storage_type_of(out)); + add_storage_type_suffix(kernel_name, graph->storage_type_of(int_weight)); + add_dtype_suffix(kernel_name, graph->dtype_of(out)); + return VK_KERNEL_FROM_STR(kernel_name); } std::string kernel_name = "linear_dq8ca_q4gsw"; @@ -1017,13 +1058,37 @@ void quantized_linear_impl( const ValueRef packed_weight_sums = prepack_standard( graph, weight_sums_data, utils::kBuffer, utils::kWidthPacked); + // The coopmat-staged-A ("-tr") dq8ca kernel reads activations through + // coopMatLoad, which can only address a row-major buffer -- so when that + // variant is active AND we know it will actually be dispatched, allocate the + // activations as kPackedInt8_4W (row-major, 4 K-values per int32) and run the + // matching packer below. dq8ca_variant_wants_rowmajor_a() is checked first so + // that on every other variant, including the shipped default, this predicate + // short-circuits to false and nothing about the stock path changes. + // + // Both halves of the decision (layout here, kernel in + // pick_linear_dqa_qw_shader) go through dq8ca_coopmat_dispatch_eligible(), so + // a shape that falls back to tiled still gets the 4h4w layout the tiled + // shader expects. + bool dq8ca_rowmajor_a = false; + if (input_quant_config.is_dynamic && weight_quant_config.nbits == 4 && + weight_quant_config.granularity == kPerGroup && + dq8ca_variant_wants_rowmajor_a()) { + dq8ca_rowmajor_a = dq8ca_coopmat_dispatch_eligible( + &graph, + output, + fp_input, + bias_data, + graph.extract_scalar(group_size)); + } + // Allocate temporary tensor to store quantized and packed input TmpTensor packed_int_input( &graph, graph.sizes_of(fp_input), vkapi::kInt8x4, utils::kBuffer, - utils::kPackedInt8_4H4W); + dq8ca_rowmajor_a ? utils::kPackedInt8_4W : utils::kPackedInt8_4H4W); // Non dynamically quantized input case if (!input_quant_config.is_dynamic) { @@ -1084,15 +1149,27 @@ void quantized_linear_impl( utils::kBuffer, utils::kWidthPacked); - add_quantize_and_pack_4h4w_with_group_sums_node( - graph, - input_quant_config, - fp_input, - int_input_sums, - packed_input_scale, - packed_input_zp, - packed_int_input, - group_size); + if (dq8ca_rowmajor_a) { + add_quantize_and_pack_4w_with_group_sums_node( + graph, + input_quant_config, + fp_input, + int_input_sums, + packed_input_scale, + packed_input_zp, + packed_int_input, + group_size); + } else { + add_quantize_and_pack_4h4w_with_group_sums_node( + graph, + input_quant_config, + fp_input, + int_input_sums, + packed_input_scale, + packed_input_zp, + packed_int_input, + group_size); + } add_linear_dqa_qw_node( graph, From a11830e91ee593a3fb08eae9910a6172c14a2cbd Mon Sep 17 00:00:00 2001 From: Yanwen Xu Date: Tue, 25 Aug 2026 15:00:35 -0700 Subject: [PATCH 11/28] [ET-VK] Cut 8da4w coopmat linear time 33.6% by hoisting the dequant correction Reduces kernel time for the 8B prefill GEMM shapes by 33.6% and raises efficiency from 28.9% to 43.5% of the architectural peak, measured on M51 (s5e9975) at pinned 980/5333/934 against main SUMD 0a88330954. That closes 58% of the gap to the reference dense-int8 vk_cooperative_matrix_perf dbuf4 shader, which reaches 54.2% on the same board, shape and clocks while doing neither int4 unpacking nor dynamic activation quantization. The starting point was a pipeline-dump ISA analysis of the shipped kernel with its group loop isolated from prologue and epilogue. That loop is 98.6% of the dynamic instruction stream, and only 4.2% of it is the matrix multiply -- a 23:1 ratio of overhead to WMMA. Two interventions moved that, one did not. The first, in linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp.glsl, applies the activation zero-point correction and the per-row activation scale once after the group loop instead of once per quantization group. The per-group term factors exactly: out = ifs * [ SUM_g wsc_g*acc_g - izp * SUM_g wsc_g*wsum_g ] The right-hand sum depends on no activation data, so it is accumulated once in the prologue from t_weight_scales and t_weight_sums -- both already bound, so this needs no new binding, no dispatch-signature change, no prepack work and no export-format change. It removes the per-group multiply and subtract, drops the wsum_sh shared array and its ping-pong, and lets izp/ifs be loaded after the loop rather than held live across it: vgpr_count 136 -> 108, v_sub_nc_u32 16 -> 0, WMMA unchanged. 18.5% faster. The second retiles from 64x32x32 to 128x64x32 with a 4x2 subgroup grid at wave32, worth a further 16.2%. Note the reference shader's own winning geometry, 128x128x64, is our *worst* valid candidate here (+31%), so that result does not transfer. Also note 128x128x64 at a 4x2 grid is not merely slow but numerically wrong -- A_ACTIVE_THREADS is 512 against a workgroup of 256, so half of A is never staged; it needs a 4x4 grid. The earlier 12/12 correctness failure at that tile was this, not register pressure. The third, in ...dbuf4zpn.glsl, widens int4 to int8 byte-parallel. The four nibbles are already one per byte, and v^8 is exactly the four-bit two's complement of v-8 because -8 == +8 (mod 16), so only a per-byte sign extension remains, done with sgn*0x1E (0x08*0x1E == 0xF0, and sgn <= 0x08080808 so it cannot carry across bytes). A naive nib-0x08080808 would borrow across byte lanes whenever a nibble is below 8. Nibble-category ops fall 92 -> 36, bit-identical output, 2.7% faster. What did not work is recorded so it is not retried: making the B shared-memory stride a power of two, to remove the integer multiplies it forces into every B address, is 4.3% slower at stride 4 and 21.8% slower at stride 8. The bank-conflict skew is worth far more than the twelve multiplies it costs. The ...dbuf4zpb.glsl variant keeps that skew templated so the measurement is reproducible. Every variant is additive and reachable only through ET_VK_DQ8CA_COOPMAT_VARIANT; the shipped default is untouched. Each passed 12 consecutive correctness runs with zero failures before any timing was taken, which this kernel family needs -- two tile defaults were shipped and reverted on 2026-08-18, one deterministically wrong and one wrong in 1 of 10 identical runs. Timings are 5 reps with clocks read back before and after, against a baseline whose spread is 0.03-0.14%. Also fixes linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr.glsl, which was committed in 7de604dcb2 failing 12/12: it bound the packed activations as a 32-bit int array while loading a coopmat from them. That mismatch is silently wrong from a StorageBuffer on this driver, though it is fine through Workgroup storage, which upstream already relies on. Binding it as int8_t gives 14/14. The dbuf4trm and dbuf4trd variants are the bisect probes that isolated it. One caveat on the whole exercise: instruction counts located the candidates but predicted none of the outcomes. The hoist beat its instruction count by 2x, the nibble change fell short of its, and the stride change went the wrong way entirely. Every claim here rests on an end-to-end measurement, not a count. --- ...near_dq8ca_q4gsw_coopmat_tsweep_dbuf4.yaml | 24 + ...ar_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr.glsl | 46 +- ...ar_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr.yaml | 47 ++ ...r_dq8ca_q4gsw_coopmat_tsweep_dbuf4trd.glsl | 682 +++++++++++++++++ ...r_dq8ca_q4gsw_coopmat_tsweep_dbuf4trd.yaml | 188 +++++ ...r_dq8ca_q4gsw_coopmat_tsweep_dbuf4trm.glsl | 707 +++++++++++++++++ ...r_dq8ca_q4gsw_coopmat_tsweep_dbuf4trm.yaml | 188 +++++ ...ar_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp.glsl | 687 +++++++++++++++++ ...ar_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp.yaml | 187 +++++ ...r_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpb.glsl | 714 ++++++++++++++++++ ...r_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpb.yaml | 42 ++ ...r_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpn.glsl | 706 +++++++++++++++++ ...r_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpn.yaml | 187 +++++ .../graph/ops/impl/QuantizedLinear.cpp | 16 +- 14 files changed, 4405 insertions(+), 16 deletions(-) create mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trd.glsl create mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trd.yaml create mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trm.glsl create mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trm.yaml create mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp.glsl create mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp.yaml create mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpb.glsl create mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpb.yaml create mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpn.glsl create mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpn.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4.yaml b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4.yaml index 54eba1e9c4e..299ba97977d 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4.yaml @@ -2162,3 +2162,27 @@ linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4: SG_GRID_X: 1 SG_GRID_Y: 8 SUBGROUP_SIZE: 32 + # Reference-matching tile: shmem_double_buf4-tr.comp's own geometry + # (TILE 128x128x16, WORKGROUP_WIDTH_IN_SUBGROUPS=4 x DBUF4_GRID_HEIGHT=2, + # wave32) => 8 subgroups / 256 threads, MMAS_PER_SG 4x2. Added so our + # kernel can be A/B'd against the teammate's PAL capture at the same + # tile and the same 2048x1024x4096 shape (8b wk_wv prefill). + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k16g42s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k16g42s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr.glsl b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr.glsl index 6027412c898..102dea7bd07 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr.glsl @@ -35,9 +35,8 @@ * coopMatLoad. (ColumnMajor is out on contiguity too: a uint packs 4 * K-values, not 4 M-values.) * - * So this shader binds t_packed_int8_input as a SCALAR int array in the - * kPackedInt8_4W layout -- plain row-major int8, 4 K-values per int32, - * row stride K4 -- produced by quantize_and_pack_4w_with_group_sums.glsl. + * So this shader binds t_packed_int8_input as a SCALAR int8_t array in the + * kPackedInt8_4W layout -- plain row-major int8, row stride K -- produced by quantize_and_pack_4w_with_group_sums.glsl. * QuantizedLinear.cpp allocates that layout (and dispatches that packer) * only when the active dq8ca variant is a "tsweep_dbuf4tr_t..." token AND * the coopmat gate passes, so the tiled fallback never sees the wrong @@ -85,6 +84,11 @@ #extension GL_KHR_shader_subgroup_basic : enable #extension GL_EXT_shader_explicit_arithmetic_types : require #extension GL_EXT_shader_explicit_arithmetic_types_int8 : require +// 8-bit SSBO access: A is bound as a scalar int8_t array so that the +// coopMatLoad below has a MATCHING component type. Loading a +// coopmat from a 32-bit int[] SSBO is what broke the first +// attempt (see header). +#extension GL_EXT_shader_8bit_storage : require #extension GL_EXT_shader_explicit_arithmetic_types_float16 : require #extension GL_EXT_control_flow_attributes : enable @@ -115,11 +119,19 @@ ${layout_declare_tensor(B, "w", "t_output", "half", IO_STORAGE, is_ // t_packed_int8_input -- but stays declared so the binding layout matches the // dispatch site. It tracks IO_STORAGE so the two IO tensors stay consistent. ${layout_declare_tensor(B, "r", "t_input", "half", IO_STORAGE, is_scalar_array=False)} -// ROW-MAJOR (kPackedInt8_4W) packed activations: scalar int array, each -// element holding 4 K-contiguous int8, row stride K4 = K/4. This is the one -// binding that differs from dbuf4 (which takes the 4h4w ivec4 block layout); -// it is what makes the coopMatLoad-based A staging below addressable. -${layout_declare_tensor(B, "r", "t_packed_int8_input", "int", "buffer", is_scalar_array=True)} +// ROW-MAJOR (kPackedInt8_4W) packed activations, bound as a scalar int8_t +// array (row stride = K int8). Two things differ from dbuf4, which takes the +// 4h4w ivec4 block layout: +// 1. row-major, so a coopMatLoad can address it at all; +// 2. element type int8_t, MATCHING the coopmat component type. +// (2) is not cosmetic. Binding the same memory as int[] and loading a +// coopmat from it -- a type mismatch that demonstrably works for the +// Workgroup storage class, which is how the MMA loop reads Ash_int8 below -- +// silently produces wrong results from a StorageBuffer on this driver. +// The reference shmem_double_buf4-tr.comp sidesteps it the same way: its +// buffer_reference is declared `A_TYPE x[]`, i.e. int8_t for the int8 config. +// All A offsets/strides here are therefore in INT8 elements, not int. +${layout_declare_tensor(B, "r", "t_packed_int8_input", "int8", "buffer", is_scalar_array=True)} ${layout_declare_tensor(B, "r", "t_int8_input_sums", "int", "buffer", is_scalar_array=True)} ${layout_declare_tensor(B, "r", "t_int8_input_scales", "half", "texture3d")} ${layout_declare_tensor(B, "r", "t_int8_input_zps", "int8", "texture3d")} @@ -217,6 +229,10 @@ void main() { const uint N = uint(output_sizes.x); const uint N4 = (N + 3u) / 4u; const uint nblocks_x_A = (K + 3u) >> 2u; + // A row stride in INT8 elements (the binding's element type). Derived from + // nblocks_x_A rather than K directly so it matches the packer's + // `m_row * K4 + k4` addressing exactly; K %% 4 == 0 makes them equal. + const uint a_row_stride_i8 = nblocks_x_A * 4u; #ifdef WEIGHT_INT4 const uint num_groups = uint(num_groups_arg); @@ -334,12 +350,12 @@ void main() { if (t < NUM_A_TILES) { const uint tm = t / A_TILES_K; const uint tk = t % A_TILES_K; - // Offset and stride are in ARRAY ELEMENTS (int = 4 int8), matching how - // the MMA loop already addresses the shared uint arrays below. + // Offset and stride are in ARRAY ELEMENTS, which for this binding are + // int8 -- i.e. the natural row-major coordinates. coopMatLoad( temp_A[s], t_packed_int8_input, - (tile_m_start + tm * MMA_M) * nblocks_x_A + tk * (MMA_K >> 2u), - nblocks_x_A, + (tile_m_start + tm * MMA_M) * a_row_stride_i8 + tk * MMA_K, + a_row_stride_i8, gl_CooperativeMatrixLayoutRowMajor); } } @@ -450,9 +466,9 @@ void main() { const uint tk = t % A_TILES_K; coopMatLoad( temp_A[s], t_packed_int8_input, - (tile_m_start + tm * MMA_M) * nblocks_x_A + - (chunkK_nxt >> 2u) + tk * (MMA_K >> 2u), - nblocks_x_A, + (tile_m_start + tm * MMA_M) * a_row_stride_i8 + chunkK_nxt + + tk * MMA_K, + a_row_stride_i8, gl_CooperativeMatrixLayoutRowMajor); } } diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr.yaml b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr.yaml index e7ee2bf9d9e..f3f6c1d7773 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr.yaml @@ -139,3 +139,50 @@ linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr: SG_GRID_X: 2 SG_GRID_Y: 4 SUBGROUP_SIZE: 32 + # Reference-matching tile: shmem_double_buf4-tr.comp's own geometry + # (TILE 128x128x16, WORKGROUP_WIDTH_IN_SUBGROUPS=4 x DBUF4_GRID_HEIGHT=2, + # wave32) => 8 subgroups / 256 threads, MMAS_PER_SG 4x2. Added so our + # kernel can be A/B'd against the teammate's PAL capture at the same + # tile and the same 2048x1024x4096 shape (8b wk_wv prefill). + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr_t128x128k16g42s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr_t128x128k16g42s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + # G2/Pavan's documented INT8 uBM config: TILE 128x128x64 with a 4x2 + # subgroup grid at wave32 (workgroupSize=256). Added to compare our kernel + # against their reference dbuf4 number (~1973us) at the same geometry and + # the same shape (M2048 K4096 N1024 == 8b wk_wv prefill). + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr_t128x128k64g42s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr_t128x128k64g42s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trd.glsl b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trd.glsl new file mode 100644 index 00000000000..2ecc26ce0f9 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trd.glsl @@ -0,0 +1,682 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +/* + * DIAGNOSTIC BISECT variant #2 of linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr. + * + * A does not go through LDS at all: each subgroup coopMatLoads the A tiles it + * needs straight from the row-major global buffer inside the MMA loop. This + * isolates coopMatLoad-from-StorageBuffer-int[] from + * coopMatStore-into-Workgroup-uint[]: dbuf4trm (manual staging) already proved + * the packer and layout correct, so if this variant passes the load is fine + * and the store is the culprit, and if it fails the load is. + * + * Ash_int8 is left allocated but unused (keeps the diff minimal; costs LDS). + * + * Original dbuf4tr header follows. + * + * "-tr" (coopmat-staged A) variant of linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4. + * + * Ported from shmem_double_buf4-tr.comp on vk_cooperative_matrix_perf's + * gemm-ubm branch. That reference file's delta over shmem_double_buf4.comp is + * that the global -> LDS staging goes through COOPERATIVE MATRIX REGISTERS + * (coopMatLoad from global -> coopmat<> array -> coopMatStore into shared) + * instead of a hand-rolled per-thread uvec4 copy, and that B lands in LDS + * column-major. + * + * Only the A half of that idea is portable to this kernel: + * + * - B is ALREADY column-major in LDS here (Bsh_int8 is K-contiguous per + * output column, read back with gl_CooperativeMatrixLayoutColumnMajor), + * so the reference's "tr" property is not a delta for B at all. + * - B CANNOT be coopmat-staged: the weights are int4, each ivec4 holding + * 8 columns x 4 K-values that need the nibble-extract / -8 / sign-pack + * below. coopMatLoad cannot unpack nibbles, and a coopmat's per-lane + * layout is opaque so one cannot be built from unpacked registers. + * B staging is therefore left byte-identical to dbuf4. + * - A CAN be coopmat-staged, but only against a ROW-MAJOR packed int8 + * activation buffer. The stock 4h4w layout (kPackedInt8_4H4W, produced by + * quantize_and_pack_4h4w_with_group_sums.glsl) is NOT row-major: element + * [m4 * K4 + k4] is an ivec4 whose COMPONENT selects one of 4 rows, so as + * a uint array the index is m4*(4*K4) + k4*4 + r, which is not affine in + * the row index and cannot be addressed by any RowMajor/ColumnMajor + * coopMatLoad. (ColumnMajor is out on contiguity too: a uint packs 4 + * K-values, not 4 M-values.) + * + * So this shader binds t_packed_int8_input as a SCALAR int array in the + * kPackedInt8_4W layout -- plain row-major int8, 4 K-values per int32, + * row stride K4 -- produced by quantize_and_pack_4w_with_group_sums.glsl. + * QuantizedLinear.cpp allocates that layout (and dispatches that packer) + * only when the active dq8ca variant is a "tsweep_dbuf4tr_t..." token AND + * the coopmat gate passes, so the tiled fallback never sees the wrong + * layout. Everything downstream of A staging -- LDS layout, int8 WMMA thread + * maps, group epilog, bias/store epilogue -- is unchanged from dbuf4. + * + * A staging (the actual -tr port): + * dbuf4: per-thread (m4, k4) ivec4 fetch; only A_ACTIVE_THREADS = + * (WG_TILE_M/4) * (WG_TILE_K/4) invocations participate, each + * scattering 4 rows into Ash_int8 with 4 scalar stores. + * dbuf4tr: per-SUBGROUP MMA_M x MMA_K tile fetch via coopMatLoad straight + * from global, then coopMatStore into the same Ash_int8 slot. + * The (WG_TILE_M/MMA_M) * (WG_TILE_K/MMA_K) tiles of a chunk are + * dealt round-robin across the NUM_SUBGROUPS subgroups. + * + * The loop structure is dbuf4's, unchanged: + * prologue: prefetch chunk 0 -> temp, store to slice 0 (no barrier) + * per iter: barrier -> prefetch(next) -> MMA(cur) -> store(next) + * and the nested `groups x chunks` loop with an unconditional group epilog is + * kept as-is (flattening it crashes the Xclipse PAL compiler at large + * spec-resolved trip counts -- see dbuf2's header). + * + * Selected at dispatch via + * ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4tr_txkgs<32|64> + * (QuantizedLinear.cpp), additive to the tsweep_dbuf4_t... and tsweep_t... + * namespaces. NOT the default -- unvalidated until it passes repeated + * test_llama_microbench --correctness-only runs (see dq8ca_coopmat_variant()'s + * comment on why a single pass is not proof). + * + * Performs: out[M,N] = dequant(int8_act) * dequant(int_w) (+ bias) + * via coopmat x coopmat -> coopmat on the matrix unit. + * + * Hard preconditions (in addition to dbuf4's): + * M % WG_TILE_M == 0, N % WG_TILE_N == 0, K % WG_TILE_K == 0, + * group_size % WG_TILE_K == 0, K % 4 == 0, + * WG_TILE_M % MMA_M == 0, WG_TILE_K % MMA_K == 0, + * t_packed_int8_input in kPackedInt8_4W (row-major) layout, + * device exposes coopmatx-> at 16x16x16. + */ + +#version 450 core + +#extension GL_KHR_cooperative_matrix : require +#extension GL_KHR_memory_scope_semantics : require +#extension GL_KHR_shader_subgroup_basic : enable +#extension GL_EXT_shader_explicit_arithmetic_types : require +#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require +// 8-bit SSBO access: A is bound as a scalar int8_t array so that the +// coopMatLoad below has a MATCHING component type. Loading a +// coopmat from a 32-bit int[] SSBO is what broke the first +// attempt (see header). +#extension GL_EXT_shader_8bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_EXT_control_flow_attributes : enable + +#define PRECISION ${PRECISION} + +$if WEIGHT_NBITS == 4: + #define WEIGHT_INT4 + +$if HAS_BIAS: + #define HAS_BIAS + +$if WEIGHT_STORAGE == "buffer": + #define WEIGHT_BUFFER + +$if IO_STORAGE == "texture3d": + #define IO_TEXTURE + +layout(std430) buffer; + +#include "common.glslh" + +// Bindings — match add_linear_dqa_qw_node arg order: +// output(0), fp_input(1), packed_int8_input(2), int_input_sums(3 - unused), +// input_scales(4), input_zps(5), packed_weight(6), weight_sums(7), +// weight_scales(8), bias(9). +${layout_declare_tensor(B, "w", "t_output", "half", IO_STORAGE, is_scalar_array=True)} +// t_input is unread here -- the activations arrive already quantized in +// t_packed_int8_input -- but stays declared so the binding layout matches the +// dispatch site. It tracks IO_STORAGE so the two IO tensors stay consistent. +${layout_declare_tensor(B, "r", "t_input", "half", IO_STORAGE, is_scalar_array=False)} +// ROW-MAJOR (kPackedInt8_4W) packed activations, bound as a scalar int8_t +// array (row stride = K int8). Two things differ from dbuf4, which takes the +// 4h4w ivec4 block layout: +// 1. row-major, so a coopMatLoad can address it at all; +// 2. element type int8_t, MATCHING the coopmat component type. +// (2) is not cosmetic. Binding the same memory as int[] and loading a +// coopmat from it -- a type mismatch that demonstrably works for the +// Workgroup storage class, which is how the MMA loop reads Ash_int8 below -- +// silently produces wrong results from a StorageBuffer on this driver. +// The reference shmem_double_buf4-tr.comp sidesteps it the same way: its +// buffer_reference is declared `A_TYPE x[]`, i.e. int8_t for the int8 config. +// All A offsets/strides here are therefore in INT8 elements, not int. +${layout_declare_tensor(B, "r", "t_packed_int8_input", "int8", "buffer", is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_int8_input_sums", "int", "buffer", is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_int8_input_scales", "half", "texture3d")} +${layout_declare_tensor(B, "r", "t_int8_input_zps", "int8", "texture3d")} +${layout_declare_tensor(B, "r", "t_packed_weight", "int", WEIGHT_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_weight_sums", "int", "buffer", is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_weight_scales", "half", "buffer", is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_bias", "half", "buffer", is_scalar_array=True)} + +${layout_declare_ubo(B, "ivec4", "output_sizes")} +${layout_declare_ubo(B, "ivec4", "input_sizes")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +${layout_declare_spec_const(C, "int", "apply_bias", "0")} +// INT4 only; inert (0) for INT8 so the dispatcher's spec list lines up. +${layout_declare_spec_const(C, "int", "K4_per_group", "0")} +${layout_declare_spec_const(C, "int", "num_groups_arg", "0")} +${layout_declare_spec_const(C, "int", "out_N_arg", "0")} + +// Tile geometry +const uint MMA_M = ${MMA_M}; +const uint MMA_N = ${MMA_N}; +const uint MMA_K = ${MMA_K}; + +const uint WG_TILE_M = ${WG_TILE_M}; +const uint WG_TILE_N = ${WG_TILE_N}; +const uint WG_TILE_K = ${WG_TILE_K}; + +const uint SG_GRID_X = ${SG_GRID_X}; +const uint SG_GRID_Y = ${SG_GRID_Y}; +const uint SUBGROUP_SIZE = ${SUBGROUP_SIZE}; +const uint NUM_SUBGROUPS = SG_GRID_X * SG_GRID_Y; +const uint WG_SIZE = NUM_SUBGROUPS * SUBGROUP_SIZE; + +const uint SG_TILE_M = WG_TILE_M / SG_GRID_Y; +const uint SG_TILE_N = WG_TILE_N / SG_GRID_X; +const uint MMAS_PER_SG_M = SG_TILE_M / MMA_M; +const uint MMAS_PER_SG_N = SG_TILE_N / MMA_N; + +const uint A_SLAB_INT8 = WG_TILE_M * MMA_K; +const uint B_USEFUL_U32 = MMA_K / 4u; +const uint B_STRIDE_U32 = B_USEFUL_U32 + 1u; // +1 skew +const uint B_SLAB_U32 = WG_TILE_N * B_STRIDE_U32; +const uint NUM_K_SLABS = WG_TILE_K / MMA_K; + +const uint A_SLAB_U32 = A_SLAB_INT8 / 4u; +const uint A_STRIDE_U32 = MMA_K / 4u; + +// One ping-pong slice covers all K-slabs of one chunk. +const uint ASH_SLICE_U32 = NUM_K_SLABS * A_SLAB_U32; +const uint BSH_SLICE_U32 = NUM_K_SLABS * B_SLAB_U32; + +// Double-buffered MMA operand staging. +shared uint Ash_int8[2u * ASH_SLICE_U32]; +shared uint Bsh_int8[2u * BSH_SLICE_U32]; + +// Per-WG-tile-row activation params (loaded ONCE at WG start; constant +// across groups). +shared int izp_sh[WG_TILE_M]; // int32 (cast from int8 source) for broadcast +shared float ifs_sh[WG_TILE_M]; // float32 (cast from fp16 source) for broadcast + +// Per-(group, output-channel) weight params, ping-ponged by group parity. +// (For per-channel INT8 only slice 0 is ever used.) +shared int wsum_sh[2u * WG_TILE_N]; +shared float wsc_sh[2u * WG_TILE_N]; + +#ifdef HAS_BIAS +shared float bias_sh[WG_TILE_N]; +#endif + +#ifdef IO_TEXTURE +// Result staging for the imageStore epilogue, mirroring the fp16 kernel: +// SG_GRID_Y bands of MMA_M rows, each WG_TILE_N wide, row-major. A full +// WG_TILE_M x WG_TILE_N buffer would cost SG_GRID_Y/MMAS_PER_SG_M x more LDS +// and wreck occupancy. float16_t-typed because coopMatStore needs it. +const uint CSH_ROWS = SG_GRID_Y * MMA_M; +shared float16_t Csh_out[CSH_ROWS * WG_TILE_N]; +#endif + +// Running fp32 accumulator (across all groups). +coopmat + result[MMAS_PER_SG_M][MMAS_PER_SG_N]; + +// Per-group int32 MMA accumulator. +coopmat + accum_int32[MMAS_PER_SG_M][MMAS_PER_SG_N]; + +void main() { + const uvec2 tileID = uvec2(gl_WorkGroupID.xy); + const uvec2 warpInTile = uvec2( + gl_SubgroupID % SG_GRID_X, + gl_SubgroupID / SG_GRID_X); + + const uint K = uint(input_sizes.x); + const uint N = uint(output_sizes.x); + const uint N4 = (N + 3u) / 4u; + const uint nblocks_x_A = (K + 3u) >> 2u; + +#ifdef WEIGHT_INT4 + const uint num_groups = uint(num_groups_arg); + const uint CHUNKS_PER_GROUP = uint(K4_per_group) * 4u / WG_TILE_K; +#else + // Per-channel: a single quant "group" spanning all of K. The nested + // groups x chunks loop below collapses to a flat chunk loop, the wsum/wsc + // ping-pong never crosses a boundary, and the epilog runs exactly once. + const uint num_groups = 1u; + const uint CHUNKS_PER_GROUP = uint(num_groups_arg); +#endif + const uint num_chunks = num_groups * CHUNKS_PER_GROUP; + + const uint tile_m_start = WG_TILE_M * tileID.y; + const uint tile_n_start = WG_TILE_N * tileID.x; + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + result[i][j] = coopmat(0.0); + accum_int32[i][j] = coopmat(0); + } + } + + const uint K_BLOCKS_PER_CHUNK = WG_TILE_K >> 2u; + + // --- A staging tile map: one MMA_M x MMA_K coopmat tile per subgroup per + // slot. A chunk holds A_TILES_M x A_TILES_K such tiles; they are dealt + // round-robin across the NUM_SUBGROUPS subgroups, so every subgroup + // participates (dbuf4's per-thread map leaves WG_SIZE - + // A_ACTIVE_THREADS invocations idle whenever the tile is small). + // A_TILES_PER_SG rounds up, so the last slot may be partially used -- + // the `t < NUM_A_TILES` guard below is subgroup-uniform (t depends only + // on gl_SubgroupID), which is what coopmat ops require. + const uint A_TILES_M = WG_TILE_M / MMA_M; + const uint A_TILES_K = WG_TILE_K / MMA_K; // == NUM_K_SLABS + const uint NUM_A_TILES = A_TILES_M * A_TILES_K; + const uint A_TILES_PER_SG = (NUM_A_TILES + NUM_SUBGROUPS - 1u) / NUM_SUBGROUPS; + +#ifdef WEIGHT_INT4 + // --- B staging thread map: (block, col) slots; each slot extracts one + // ColumnMajor LDS uint (4 K-contiguous sign-extended int8) --- + const uint B_TOTAL_SLOTS = K_BLOCKS_PER_CHUNK * WG_TILE_N; + const uint B_SLOTS_PER_THREAD = B_TOTAL_SLOTS / WG_SIZE; + const uint N8_PER_TILE = WG_TILE_N >> 3u; +#else + // --- B staging thread map: one (k4, n4) ivec4 block per active thread --- + const uint B_FETCH_SLOTS = K_BLOCKS_PER_CHUNK * (WG_TILE_N >> 2u); + const uint N4_PER_TILE = WG_TILE_N >> 2u; + const uint b_k4_in_chunk = gl_LocalInvocationID.x / N4_PER_TILE; + const uint b_n_uint_col = gl_LocalInvocationID.x % N4_PER_TILE; + const bool b_active = gl_LocalInvocationID.x < B_FETCH_SLOTS; +#endif + + // Prefetch temp registers. temp_A is a coopmat array (the -tr change); + // indices into it are [[unroll]]-resolved compile-time constants, never + // dynamic -- dynamic indexing of a coopmat array is exactly the construct + // the Xclipse/AMD-PAL compiler has miscompiled before. + // (temp_A removed) +#ifdef WEIGHT_INT4 + ivec4 temp_B[B_SLOTS_PER_THREAD]; + int temp_wsum; + float temp_wsc; +#else + ivec4 temp_B; +#endif + + // ========================================================= + // PROLOGUE + // ========================================================= + if (gl_LocalInvocationID.x < (WG_TILE_M >> 2u)) { + const uint m4 = (tile_m_start >> 2u) + gl_LocalInvocationID.x; + const vec4 sc = vec4(texelFetch(t_int8_input_scales, ivec3(m4, 0, 0), 0)); + const ivec4 zp = texelFetch(t_int8_input_zps, ivec3(m4, 0, 0), 0); + const uint base = gl_LocalInvocationID.x * 4u; + ifs_sh[base + 0u] = sc.x; ifs_sh[base + 1u] = sc.y; + ifs_sh[base + 2u] = sc.z; ifs_sh[base + 3u] = sc.w; + izp_sh[base + 0u] = zp.x; izp_sh[base + 1u] = zp.y; + izp_sh[base + 2u] = zp.z; izp_sh[base + 3u] = zp.w; + } + // Group 0 weight sums/scales -> slice 0. + if (gl_LocalInvocationID.x < WG_TILE_N) { + const uint n_idx = tile_n_start + gl_LocalInvocationID.x; + f16vec4 sv = t_weight_scales[n_idx >> 2u]; + wsc_sh[gl_LocalInvocationID.x] = float(sv[n_idx & 3u]); + wsum_sh[gl_LocalInvocationID.x] = t_weight_sums[n_idx]; + } + memoryBarrierShared(); + barrier(); + + // izp/ifs are per-row activation params, constant across K groups — + // broadcast them into coopmats ONCE; the group epilog reuses them every + // group (they depend only on the row block i, not on the group or j). + coopmat + izp_bcast[MMAS_PER_SG_M]; + coopmat + ifs_bcast[MMAS_PER_SG_M]; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + const uint local_m_base = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + coopMatLoad( + izp_bcast[i], izp_sh, + local_m_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutColumnMajor); + coopMatLoad( + ifs_bcast[i], ifs_sh, + local_m_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutColumnMajor); + } + + // dbuf4: prefetch chunk 0 into temp registers, THEN store to slice 0 (no + // barrier here -- the main loop's first iteration barriers before + // reading slice 0). + // (A prefetch removed: loaded directly in the MMA loop) +#ifdef WEIGHT_INT4 + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { + const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; + const uint block_in_chunk = slot >> 3u; + const uint k4_blk = block_in_chunk / N8_PER_TILE; + const uint n8_blk = (tile_n_start >> 3u) + (block_in_chunk % N8_PER_TILE); +#ifdef WEIGHT_BUFFER + temp_B[si] = t_packed_weight[(n8_blk * nblocks_x_A) + k4_blk]; +#else + temp_B[si] = texelFetch(t_packed_weight, ivec2(k4_blk, n8_blk), 0); +#endif + } +#else + if (b_active) { + const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; +#ifdef WEIGHT_BUFFER + temp_B = t_packed_weight[(b_k4_in_chunk * N4) + block_x_w]; +#else + temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, b_k4_in_chunk), 0); +#endif + } +#endif + { + // store chunk 0 -> slice 0 + // (A LDS store removed) +#ifdef WEIGHT_INT4 + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { + const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; + const uint block_in_chunk = slot >> 3u; + const uint col_in_block = slot & 7u; + const uint k4_in_chunk = block_in_chunk / N8_PER_TILE; + const uint n8_in_tile = block_in_chunk % N8_PER_TILE; + const uint r = col_in_block & 3u; + const uint parity = col_in_block >> 2u; + const int w = temp_B[si][r]; + const int base = int(4u * parity); + const int v0 = (((w >> (base + 0)) & 0xF) - 8) & 0xFF; + const int v1 = (((w >> (base + 8)) & 0xF) - 8) & 0xFF; + const int v2 = (((w >> (base + 16)) & 0xF) - 8) & 0xFF; + const int v3 = (((w >> (base + 24)) & 0xF) - 8) & 0xFF; + const uint n_col = n8_in_tile * 8u + r + parity * 4u; + const uint slab_idx = k4_in_chunk / (MMA_K >> 2u); + const uint k4_in_slab = k4_in_chunk % (MMA_K >> 2u); + Bsh_int8[slab_idx * B_SLAB_U32 + n_col * B_STRIDE_U32 + k4_in_slab] = + uint(v0 | (v1 << 8) | (v2 << 16) | (v3 << 24)); + } +#else + if (b_active) { + const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); + const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); + const uint n_col_base = b_n_uint_col * 4u; + [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { + Bsh_int8[slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab] = + uint(temp_B[n_in_blk]); + } + } +#endif + } + + // ========================================================= + // MAIN LOOP (dbuf4) — nested groups x chunks (kept nested; flattening it + // with a conditional coopmat epilog crashes the Xclipse PAL compiler at + // large spec-resolved trip counts). One barrier per chunk. Chunk + // iteration (global index `chunk`): + // 1. barrier — A/B slice (chunk%2) fully written; on the first chunk + // of group g, wsum/wsc slice (g%2) is too. + // 2. prefetch — chunk+1 (A blocks, B blocks) into temp; when chunk+1 + // starts a new group, also its wsum/wsc element. Skipped + // entirely on the final chunk. + // 3. int8 MMA — on slice (chunk%2) into accum_int32. + // 4. store — temp -> A/B slice ((chunk+1)%2), unpacking the weight; + // on a group boundary, wsum/wsc -> slice ((g+1)%2). + // The group epilog runs unconditionally at the tail of each group. + // ========================================================= + uint chunk = 0; + for (uint group_i = 0; group_i < num_groups; ++group_i) { + for (uint inner = 0; inner < CHUNKS_PER_GROUP; ++inner, ++chunk) { + const bool has_next = chunk + 1u < num_chunks; + const bool group_crossing = has_next && (inner + 1u == CHUNKS_PER_GROUP); + const uint cur_a = (chunk % 2u) * ASH_SLICE_U32; + const uint cur_b = (chunk % 2u) * BSH_SLICE_U32; + const uint nxt_a = ((chunk + 1u) % 2u) * ASH_SLICE_U32; + const uint nxt_b = ((chunk + 1u) % 2u) * BSH_SLICE_U32; + + barrier(); + + // --- 2. prefetch chunk+1 -> temp --- + if (has_next) { + const uint chunkK_nxt = (chunk + 1u) * WG_TILE_K; + // (A prefetch removed) +#ifdef WEIGHT_INT4 + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { + const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; + const uint block_in_chunk = slot >> 3u; + const uint k4_blk = (chunkK_nxt >> 2u) + block_in_chunk / N8_PER_TILE; + const uint n8_blk = (tile_n_start >> 3u) + (block_in_chunk % N8_PER_TILE); +#ifdef WEIGHT_BUFFER + temp_B[si] = t_packed_weight[(n8_blk * nblocks_x_A) + k4_blk]; +#else + temp_B[si] = texelFetch(t_packed_weight, ivec2(k4_blk, n8_blk), 0); +#endif + } + if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { + const uint n_idx = tile_n_start + gl_LocalInvocationID.x; + f16vec4 sv = t_weight_scales[(group_i + 1u) * N4 + (n_idx >> 2u)]; + temp_wsc = float(sv[n_idx & 3u]); + temp_wsum = t_weight_sums[(group_i + 1u) * N + n_idx]; + } +#else + if (b_active) { + const uint block_y_w = (chunkK_nxt >> 2u) + b_k4_in_chunk; + const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; +#ifdef WEIGHT_BUFFER + temp_B = t_packed_weight[(block_y_w * N4) + block_x_w]; +#else + temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, block_y_w), 0); +#endif + } +#endif + } + + // --- 3. int8 MMA on the cur slice --- + [[unroll]] for (uint k = 0; k < NUM_K_SLABS; ++k) { + const uint slab_a_base_u32 = cur_a + k * A_SLAB_U32; + const uint slab_b_base_u32 = cur_b + k * B_SLAB_U32; + + coopmat matA[MMAS_PER_SG_M]; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + const uint row_a = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + // Offsets/strides are in int8 elements, matching the int8_t + // binding. With A bound as int[] instead, this load silently + // produced wrong results regardless of which unit was used. + const uint a_row_stride_i8 = nblocks_x_A * 4u; // int8 elements + coopMatLoad( + matA[i], t_packed_int8_input, + (tile_m_start + row_a) * a_row_stride_i8 + + (chunk * WG_TILE_K + k * MMA_K), + a_row_stride_i8, + gl_CooperativeMatrixLayoutRowMajor); + } + + coopmat matB; + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint col_b = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopMatLoad( + matB, Bsh_int8, + slab_b_base_u32 + col_b * B_STRIDE_U32, + B_STRIDE_U32, + gl_CooperativeMatrixLayoutColumnMajor); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + accum_int32[i][j] = coopMatMulAdd(matA[i], matB, accum_int32[i][j]); + } + } + } + + // --- 4. store temp (chunk+1) -> nxt slice --- + if (has_next) { + // (A LDS store removed) +#ifdef WEIGHT_INT4 + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { + const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; + const uint block_in_chunk = slot >> 3u; + const uint col_in_block = slot & 7u; + const uint k4_in_chunk = block_in_chunk / N8_PER_TILE; + const uint n8_in_tile = block_in_chunk % N8_PER_TILE; + const uint r = col_in_block & 3u; + const uint parity = col_in_block >> 2u; + const int w = temp_B[si][r]; + const int base = int(4u * parity); + const int v0 = (((w >> (base + 0)) & 0xF) - 8) & 0xFF; + const int v1 = (((w >> (base + 8)) & 0xF) - 8) & 0xFF; + const int v2 = (((w >> (base + 16)) & 0xF) - 8) & 0xFF; + const int v3 = (((w >> (base + 24)) & 0xF) - 8) & 0xFF; + const uint n_col = n8_in_tile * 8u + r + parity * 4u; + const uint slab_idx = k4_in_chunk / (MMA_K >> 2u); + const uint k4_in_slab = k4_in_chunk % (MMA_K >> 2u); + Bsh_int8[nxt_b + slab_idx * B_SLAB_U32 + n_col * B_STRIDE_U32 + k4_in_slab] = + uint(v0 | (v1 << 8) | (v2 << 16) | (v3 << 24)); + } + if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { + const uint wbase_nxt = ((group_i + 1u) % 2u) * WG_TILE_N; + wsum_sh[wbase_nxt + gl_LocalInvocationID.x] = temp_wsum; + wsc_sh[wbase_nxt + gl_LocalInvocationID.x] = temp_wsc; + } +#else + if (b_active) { + const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); + const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); + const uint n_col_base = b_n_uint_col * 4u; + [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { + Bsh_int8[nxt_b + slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab] = + uint(temp_B[n_in_blk]); + } + } +#endif + } + } // chunks + + // --- Group epilog: dequant accum_int32 -> result, reset accum --- + { + const uint wbase = (group_i % 2u) * WG_TILE_N; + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + + coopmat wsum_bcast; + coopMatLoad( + wsum_bcast, wsum_sh, + wbase + local_n_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutRowMajor); + + coopmat wsc_bcast; + coopMatLoad( + wsc_bcast, wsc_sh, + wbase + local_n_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutRowMajor); + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + coopmat adjusted = + accum_int32[i][j] - izp_bcast[i] * wsum_bcast; + coopmat adjusted_fp = + coopmat(adjusted); + coopmat scales_outer = + ifs_bcast[i] * wsc_bcast; + result[i][j] += adjusted_fp * scales_outer; + accum_int32[i][j] = coopmat(0); + } + } + } + } // groups + + // --- Bias (optional) --- +#ifdef HAS_BIAS + if (apply_bias > 0) { + for (uint t = gl_LocalInvocationID.x; t < WG_TILE_N; t += WG_SIZE) { + bias_sh[t] = float(t_bias[tile_n_start + t]); + } + memoryBarrierShared(); + barrier(); + } +#endif + + // --- Store result tile --- + // N for the store address math MUST come from the spec constant, not the + // sizes UBO (see out_N_arg above). +#ifdef IO_TEXTURE + // Epilogue iteration i drains accumulator row-block i from EVERY subgroup + // into Csh_out at once, so the SG_GRID_Y bands it holds are disjoint global + // row ranges; the whole workgroup then imageStores them. lr / MMA_M is the + // writing subgroup's warpInTile.y, so the global row reproduces the buffer + // path's gi exactly. + // + // PORTABILITY NOTE: the barrier() in the loop body keeps this loop rolled + // despite [[unroll]], so result[i][j] IS dynamically indexed. Coopmat arrays + // are opaque per-lane storage and dynamic indexing is exactly the construct + // the Xclipse/AMD-PAL compiler has broken before -- check this first if the + // texture variants miscompile on M51. + const uint CSH_TEXELS_PER_ROW = WG_TILE_N / 4u; + const uint CSH_TEXELS = CSH_ROWS * CSH_TEXELS_PER_ROW; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + // Guards Csh_out against the previous iteration's readers. Inert on i == 0 + // but must stay unconditional to remain workgroup-uniform. + barrier(); + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { +#ifdef HAS_BIAS + if (apply_bias > 0) { + const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopmat bias_tile; + coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); + result[i][j] += bias_tile; + } +#endif + coopmat out_tile = + coopmat(result[i][j]); + coopMatStore( + out_tile, Csh_out, + warpInTile.y * MMA_M * WG_TILE_N + + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j), + WG_TILE_N, + gl_CooperativeMatrixLayoutRowMajor); + } + memoryBarrierShared(); + barrier(); + + for (uint t = gl_LocalInvocationID.x; t < CSH_TEXELS; t += WG_SIZE) { + const uint lr = t / CSH_TEXELS_PER_ROW; + const uint lc4 = t % CSH_TEXELS_PER_ROW; + const uint m = + tile_m_start + (lr / MMA_M) * SG_TILE_M + i * MMA_M + (lr % MMA_M); + const uint base = lr * WG_TILE_N + lc4 * 4u; + imageStore( + t_output, + ivec3(tile_n_start / 4u + lc4, m, 0), + vec4( + float(Csh_out[base]), + float(Csh_out[base + 1u]), + float(Csh_out[base + 2u]), + float(Csh_out[base + 3u]))); + } + } +#else + const uint N_out = uint(out_N_arg); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint gi = tile_m_start + MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + const uint gj = tile_n_start + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + +#ifdef HAS_BIAS + if (apply_bias > 0) { + const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopmat bias_tile; + coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); + result[i][j] += bias_tile; + } +#endif + + coopmat out_tile = + coopmat(result[i][j]); + coopMatStore( + out_tile, t_output, + gi * N_out + gj, N_out, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#endif // IO_TEXTURE +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trd.yaml b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trd.yaml new file mode 100644 index 00000000000..e26c6c3bea6 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trd.yaml @@ -0,0 +1,188 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# "-tr" (coopmat-staged A) variants of the int8 dq8ca_q4gsw coopmat kernel, +# ported from shmem_double_buf4-tr.comp (vk_cooperative_matrix_perf, +# gemm-ubm). linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trd.glsl is a fork of +# linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4.glsl with ONLY the A-side global -> +# LDS staging swapped from a per-thread ivec4 copy to +# coopMatLoad(global) -> coopmat<> -> coopMatStore(LDS). B staging is +# byte-identical (int4 nibble unpack cannot be coopmat-staged) and B was +# already column-major in LDS, so it is not a delta. +# +# These variants require t_packed_int8_input in the ROW-MAJOR +# kPackedInt8_4W layout (produced by quantize_and_pack_4w_with_group_sums), +# NOT the stock 4h4w block layout. QuantizedLinear.cpp switches the tensor +# layout and the packer node together with the variant token, and only when +# the coopmat gate passes -- see dq8ca_wants_rowmajor_int8_input(). +# +# Selected via +# ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4trd_txkgs. +# NOT the shipped default: unvalidated until it passes repeated +# test_llama_microbench --correctness-only runs (a single pass is not proof -- +# see dq8ca_coopmat_variant()'s comment). +# +# Tile-geometry preconditions beyond dbuf4's: WG_TILE_M % MMA_M == 0 and +# WG_TILE_K % MMA_K == 0 (both hold for every entry below). The seed set is +# deliberately small -- widen it from a sweep once correctness is established. + +linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trd: + parameter_names_with_default_values: + PRECISION: highp + HAS_BIAS: false + IO_STORAGE: buffer + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + MMA_M: 16 + MMA_N: 16 + MMA_K: 16 + WG_TILE_M: 64 + WG_TILE_N: 32 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + shader_variants: + # production 8da4w tile (dbuf2/dbuf4 default) -- the A/B anchor + # A tiles/chunk = 8, subgroups = 2 -> A_TILES_PER_SG = 4 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trd_t64x32k32g12s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 32 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trd_t64x32k32g12s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 32 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trd_t64x32k32g12s64_buffer_buffer_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: buffer + WG_TILE_M: 64 + WG_TILE_N: 32 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + # q4gsw production geometry; 8 A-tiles over 4 subgroups + # A tiles/chunk = 8, subgroups = 4 -> A_TILES_PER_SG = 2 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trd_t128x128k16g22s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trd_t128x128k16g22s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + # 1 A-tile per subgroup (A_TILES_PER_SG == 1, no leftover slot) + # A tiles/chunk = 8, subgroups = 8 -> A_TILES_PER_SG = 1 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trd_t64x64k32g24s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trd_t64x64k32g24s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + # 16 A-tiles over 8 subgroups + # A tiles/chunk = 16, subgroups = 8 -> A_TILES_PER_SG = 2 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trd_t128x64k32g24s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trd_t128x64k32g24s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + # Reference-matching tile: shmem_double_buf4-tr.comp's own geometry + # (TILE 128x128x16, WORKGROUP_WIDTH_IN_SUBGROUPS=4 x DBUF4_GRID_HEIGHT=2, + # wave32) => 8 subgroups / 256 threads, MMAS_PER_SG 4x2. Added so our + # kernel can be A/B'd against the teammate's PAL capture at the same + # tile and the same 2048x1024x4096 shape (8b wk_wv prefill). + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trd_t128x128k16g42s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trd_t128x128k16g42s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + # G2/Pavan's documented INT8 uBM config: TILE 128x128x64 with a 4x2 + # subgroup grid at wave32 (workgroupSize=256). Added to compare our kernel + # against their reference dbuf4 number (~1973us) at the same geometry and + # the same shape (M2048 K4096 N1024 == 8b wk_wv prefill). + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trd_t128x128k64g42s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trd_t128x128k64g42s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trm.glsl b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trm.glsl new file mode 100644 index 00000000000..b12d6c21012 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trm.glsl @@ -0,0 +1,707 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +/* + * DIAGNOSTIC BISECT variant of linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr. + * + * Identical to dbuf4tr -- same ROW-MAJOR kPackedInt8_4W activation layout, same + * quantize_and_pack_4w_with_group_sums producer, same LDS layout, same loop -- + * EXCEPT that A staging is done with plain scalar loads/stores instead of + * coopMatLoad/coopMatStore. It exists to split a dbuf4tr correctness failure + * into "the row-major packer/layout is wrong" (this variant also fails) vs + * "the coopmat staging is wrong" (this variant passes). + * + * Not a perf candidate; delete once dbuf4tr is understood. + * + * Original dbuf4tr header follows. + * + * "-tr" (coopmat-staged A) variant of linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4. + * + * Ported from shmem_double_buf4-tr.comp on vk_cooperative_matrix_perf's + * gemm-ubm branch. That reference file's delta over shmem_double_buf4.comp is + * that the global -> LDS staging goes through COOPERATIVE MATRIX REGISTERS + * (coopMatLoad from global -> coopmat<> array -> coopMatStore into shared) + * instead of a hand-rolled per-thread uvec4 copy, and that B lands in LDS + * column-major. + * + * Only the A half of that idea is portable to this kernel: + * + * - B is ALREADY column-major in LDS here (Bsh_int8 is K-contiguous per + * output column, read back with gl_CooperativeMatrixLayoutColumnMajor), + * so the reference's "tr" property is not a delta for B at all. + * - B CANNOT be coopmat-staged: the weights are int4, each ivec4 holding + * 8 columns x 4 K-values that need the nibble-extract / -8 / sign-pack + * below. coopMatLoad cannot unpack nibbles, and a coopmat's per-lane + * layout is opaque so one cannot be built from unpacked registers. + * B staging is therefore left byte-identical to dbuf4. + * - A CAN be coopmat-staged, but only against a ROW-MAJOR packed int8 + * activation buffer. The stock 4h4w layout (kPackedInt8_4H4W, produced by + * quantize_and_pack_4h4w_with_group_sums.glsl) is NOT row-major: element + * [m4 * K4 + k4] is an ivec4 whose COMPONENT selects one of 4 rows, so as + * a uint array the index is m4*(4*K4) + k4*4 + r, which is not affine in + * the row index and cannot be addressed by any RowMajor/ColumnMajor + * coopMatLoad. (ColumnMajor is out on contiguity too: a uint packs 4 + * K-values, not 4 M-values.) + * + * So this shader binds t_packed_int8_input as a SCALAR int array in the + * kPackedInt8_4W layout -- plain row-major int8, 4 K-values per int32, + * row stride K4 -- produced by quantize_and_pack_4w_with_group_sums.glsl. + * QuantizedLinear.cpp allocates that layout (and dispatches that packer) + * only when the active dq8ca variant is a "tsweep_dbuf4tr_t..." token AND + * the coopmat gate passes, so the tiled fallback never sees the wrong + * layout. Everything downstream of A staging -- LDS layout, int8 WMMA thread + * maps, group epilog, bias/store epilogue -- is unchanged from dbuf4. + * + * A staging (the actual -tr port): + * dbuf4: per-thread (m4, k4) ivec4 fetch; only A_ACTIVE_THREADS = + * (WG_TILE_M/4) * (WG_TILE_K/4) invocations participate, each + * scattering 4 rows into Ash_int8 with 4 scalar stores. + * dbuf4tr: per-SUBGROUP MMA_M x MMA_K tile fetch via coopMatLoad straight + * from global, then coopMatStore into the same Ash_int8 slot. + * The (WG_TILE_M/MMA_M) * (WG_TILE_K/MMA_K) tiles of a chunk are + * dealt round-robin across the NUM_SUBGROUPS subgroups. + * + * The loop structure is dbuf4's, unchanged: + * prologue: prefetch chunk 0 -> temp, store to slice 0 (no barrier) + * per iter: barrier -> prefetch(next) -> MMA(cur) -> store(next) + * and the nested `groups x chunks` loop with an unconditional group epilog is + * kept as-is (flattening it crashes the Xclipse PAL compiler at large + * spec-resolved trip counts -- see dbuf2's header). + * + * Selected at dispatch via + * ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4tr_txkgs<32|64> + * (QuantizedLinear.cpp), additive to the tsweep_dbuf4_t... and tsweep_t... + * namespaces. NOT the default -- unvalidated until it passes repeated + * test_llama_microbench --correctness-only runs (see dq8ca_coopmat_variant()'s + * comment on why a single pass is not proof). + * + * Performs: out[M,N] = dequant(int8_act) * dequant(int_w) (+ bias) + * via coopmat x coopmat -> coopmat on the matrix unit. + * + * Hard preconditions (in addition to dbuf4's): + * M % WG_TILE_M == 0, N % WG_TILE_N == 0, K % WG_TILE_K == 0, + * group_size % WG_TILE_K == 0, K % 4 == 0, + * WG_TILE_M % MMA_M == 0, WG_TILE_K % MMA_K == 0, + * t_packed_int8_input in kPackedInt8_4W (row-major) layout, + * device exposes coopmatx-> at 16x16x16. + */ + +#version 450 core + +#extension GL_KHR_cooperative_matrix : require +#extension GL_KHR_memory_scope_semantics : require +#extension GL_KHR_shader_subgroup_basic : enable +#extension GL_EXT_shader_explicit_arithmetic_types : require +#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_EXT_control_flow_attributes : enable + +#define PRECISION ${PRECISION} + +$if WEIGHT_NBITS == 4: + #define WEIGHT_INT4 + +$if HAS_BIAS: + #define HAS_BIAS + +$if WEIGHT_STORAGE == "buffer": + #define WEIGHT_BUFFER + +$if IO_STORAGE == "texture3d": + #define IO_TEXTURE + +layout(std430) buffer; + +#include "common.glslh" + +// Bindings — match add_linear_dqa_qw_node arg order: +// output(0), fp_input(1), packed_int8_input(2), int_input_sums(3 - unused), +// input_scales(4), input_zps(5), packed_weight(6), weight_sums(7), +// weight_scales(8), bias(9). +${layout_declare_tensor(B, "w", "t_output", "half", IO_STORAGE, is_scalar_array=True)} +// t_input is unread here -- the activations arrive already quantized in +// t_packed_int8_input -- but stays declared so the binding layout matches the +// dispatch site. It tracks IO_STORAGE so the two IO tensors stay consistent. +${layout_declare_tensor(B, "r", "t_input", "half", IO_STORAGE, is_scalar_array=False)} +// ROW-MAJOR (kPackedInt8_4W) packed activations: scalar int array, each +// element holding 4 K-contiguous int8, row stride K4 = K/4. This is the one +// binding that differs from dbuf4 (which takes the 4h4w ivec4 block layout); +// it is what makes the coopMatLoad-based A staging below addressable. +${layout_declare_tensor(B, "r", "t_packed_int8_input", "int", "buffer", is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_int8_input_sums", "int", "buffer", is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_int8_input_scales", "half", "texture3d")} +${layout_declare_tensor(B, "r", "t_int8_input_zps", "int8", "texture3d")} +${layout_declare_tensor(B, "r", "t_packed_weight", "int", WEIGHT_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_weight_sums", "int", "buffer", is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_weight_scales", "half", "buffer", is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_bias", "half", "buffer", is_scalar_array=True)} + +${layout_declare_ubo(B, "ivec4", "output_sizes")} +${layout_declare_ubo(B, "ivec4", "input_sizes")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +${layout_declare_spec_const(C, "int", "apply_bias", "0")} +// INT4 only; inert (0) for INT8 so the dispatcher's spec list lines up. +${layout_declare_spec_const(C, "int", "K4_per_group", "0")} +${layout_declare_spec_const(C, "int", "num_groups_arg", "0")} +${layout_declare_spec_const(C, "int", "out_N_arg", "0")} + +// Tile geometry +const uint MMA_M = ${MMA_M}; +const uint MMA_N = ${MMA_N}; +const uint MMA_K = ${MMA_K}; + +const uint WG_TILE_M = ${WG_TILE_M}; +const uint WG_TILE_N = ${WG_TILE_N}; +const uint WG_TILE_K = ${WG_TILE_K}; + +const uint SG_GRID_X = ${SG_GRID_X}; +const uint SG_GRID_Y = ${SG_GRID_Y}; +const uint SUBGROUP_SIZE = ${SUBGROUP_SIZE}; +const uint NUM_SUBGROUPS = SG_GRID_X * SG_GRID_Y; +const uint WG_SIZE = NUM_SUBGROUPS * SUBGROUP_SIZE; + +const uint SG_TILE_M = WG_TILE_M / SG_GRID_Y; +const uint SG_TILE_N = WG_TILE_N / SG_GRID_X; +const uint MMAS_PER_SG_M = SG_TILE_M / MMA_M; +const uint MMAS_PER_SG_N = SG_TILE_N / MMA_N; + +const uint A_SLAB_INT8 = WG_TILE_M * MMA_K; +const uint B_USEFUL_U32 = MMA_K / 4u; +const uint B_STRIDE_U32 = B_USEFUL_U32 + 1u; // +1 skew +const uint B_SLAB_U32 = WG_TILE_N * B_STRIDE_U32; +const uint NUM_K_SLABS = WG_TILE_K / MMA_K; + +const uint A_SLAB_U32 = A_SLAB_INT8 / 4u; +const uint A_STRIDE_U32 = MMA_K / 4u; + +// One ping-pong slice covers all K-slabs of one chunk. +const uint ASH_SLICE_U32 = NUM_K_SLABS * A_SLAB_U32; +const uint BSH_SLICE_U32 = NUM_K_SLABS * B_SLAB_U32; + +// Double-buffered MMA operand staging. +shared uint Ash_int8[2u * ASH_SLICE_U32]; +shared uint Bsh_int8[2u * BSH_SLICE_U32]; + +// Per-WG-tile-row activation params (loaded ONCE at WG start; constant +// across groups). +shared int izp_sh[WG_TILE_M]; // int32 (cast from int8 source) for broadcast +shared float ifs_sh[WG_TILE_M]; // float32 (cast from fp16 source) for broadcast + +// Per-(group, output-channel) weight params, ping-ponged by group parity. +// (For per-channel INT8 only slice 0 is ever used.) +shared int wsum_sh[2u * WG_TILE_N]; +shared float wsc_sh[2u * WG_TILE_N]; + +#ifdef HAS_BIAS +shared float bias_sh[WG_TILE_N]; +#endif + +#ifdef IO_TEXTURE +// Result staging for the imageStore epilogue, mirroring the fp16 kernel: +// SG_GRID_Y bands of MMA_M rows, each WG_TILE_N wide, row-major. A full +// WG_TILE_M x WG_TILE_N buffer would cost SG_GRID_Y/MMAS_PER_SG_M x more LDS +// and wreck occupancy. float16_t-typed because coopMatStore needs it. +const uint CSH_ROWS = SG_GRID_Y * MMA_M; +shared float16_t Csh_out[CSH_ROWS * WG_TILE_N]; +#endif + +// Running fp32 accumulator (across all groups). +coopmat + result[MMAS_PER_SG_M][MMAS_PER_SG_N]; + +// Per-group int32 MMA accumulator. +coopmat + accum_int32[MMAS_PER_SG_M][MMAS_PER_SG_N]; + +void main() { + const uvec2 tileID = uvec2(gl_WorkGroupID.xy); + const uvec2 warpInTile = uvec2( + gl_SubgroupID % SG_GRID_X, + gl_SubgroupID / SG_GRID_X); + + const uint K = uint(input_sizes.x); + const uint N = uint(output_sizes.x); + const uint N4 = (N + 3u) / 4u; + const uint nblocks_x_A = (K + 3u) >> 2u; + +#ifdef WEIGHT_INT4 + const uint num_groups = uint(num_groups_arg); + const uint CHUNKS_PER_GROUP = uint(K4_per_group) * 4u / WG_TILE_K; +#else + // Per-channel: a single quant "group" spanning all of K. The nested + // groups x chunks loop below collapses to a flat chunk loop, the wsum/wsc + // ping-pong never crosses a boundary, and the epilog runs exactly once. + const uint num_groups = 1u; + const uint CHUNKS_PER_GROUP = uint(num_groups_arg); +#endif + const uint num_chunks = num_groups * CHUNKS_PER_GROUP; + + const uint tile_m_start = WG_TILE_M * tileID.y; + const uint tile_n_start = WG_TILE_N * tileID.x; + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + result[i][j] = coopmat(0.0); + accum_int32[i][j] = coopmat(0); + } + } + + const uint K_BLOCKS_PER_CHUNK = WG_TILE_K >> 2u; + + // --- A staging slot map (DIAGNOSTIC): one row-major int (4 K-contiguous + // int8 of ONE row) per thread per slot. Deliberately NOT coopmat. + const uint A_TOTAL_SLOTS = WG_TILE_M * K_BLOCKS_PER_CHUNK; + const uint A_SLOTS_PER_THREAD = (A_TOTAL_SLOTS + WG_SIZE - 1u) / WG_SIZE; + + // --- (unused in this variant) A staging tile map: one MMA_M x MMA_K coopmat tile per subgroup per + // slot. A chunk holds A_TILES_M x A_TILES_K such tiles; they are dealt + // round-robin across the NUM_SUBGROUPS subgroups, so every subgroup + // participates (dbuf4's per-thread map leaves WG_SIZE - + // A_ACTIVE_THREADS invocations idle whenever the tile is small). + // A_TILES_PER_SG rounds up, so the last slot may be partially used -- + // the `t < NUM_A_TILES` guard below is subgroup-uniform (t depends only + // on gl_SubgroupID), which is what coopmat ops require. + const uint A_TILES_M = WG_TILE_M / MMA_M; + const uint A_TILES_K = WG_TILE_K / MMA_K; // == NUM_K_SLABS + const uint NUM_A_TILES = A_TILES_M * A_TILES_K; + const uint A_TILES_PER_SG = (NUM_A_TILES + NUM_SUBGROUPS - 1u) / NUM_SUBGROUPS; + +#ifdef WEIGHT_INT4 + // --- B staging thread map: (block, col) slots; each slot extracts one + // ColumnMajor LDS uint (4 K-contiguous sign-extended int8) --- + const uint B_TOTAL_SLOTS = K_BLOCKS_PER_CHUNK * WG_TILE_N; + const uint B_SLOTS_PER_THREAD = B_TOTAL_SLOTS / WG_SIZE; + const uint N8_PER_TILE = WG_TILE_N >> 3u; +#else + // --- B staging thread map: one (k4, n4) ivec4 block per active thread --- + const uint B_FETCH_SLOTS = K_BLOCKS_PER_CHUNK * (WG_TILE_N >> 2u); + const uint N4_PER_TILE = WG_TILE_N >> 2u; + const uint b_k4_in_chunk = gl_LocalInvocationID.x / N4_PER_TILE; + const uint b_n_uint_col = gl_LocalInvocationID.x % N4_PER_TILE; + const bool b_active = gl_LocalInvocationID.x < B_FETCH_SLOTS; +#endif + + // Prefetch temp registers. temp_A is a coopmat array (the -tr change); + // indices into it are [[unroll]]-resolved compile-time constants, never + // dynamic -- dynamic indexing of a coopmat array is exactly the construct + // the Xclipse/AMD-PAL compiler has miscompiled before. + int temp_A[A_SLOTS_PER_THREAD]; +#ifdef WEIGHT_INT4 + ivec4 temp_B[B_SLOTS_PER_THREAD]; + int temp_wsum; + float temp_wsc; +#else + ivec4 temp_B; +#endif + + // ========================================================= + // PROLOGUE + // ========================================================= + if (gl_LocalInvocationID.x < (WG_TILE_M >> 2u)) { + const uint m4 = (tile_m_start >> 2u) + gl_LocalInvocationID.x; + const vec4 sc = vec4(texelFetch(t_int8_input_scales, ivec3(m4, 0, 0), 0)); + const ivec4 zp = texelFetch(t_int8_input_zps, ivec3(m4, 0, 0), 0); + const uint base = gl_LocalInvocationID.x * 4u; + ifs_sh[base + 0u] = sc.x; ifs_sh[base + 1u] = sc.y; + ifs_sh[base + 2u] = sc.z; ifs_sh[base + 3u] = sc.w; + izp_sh[base + 0u] = zp.x; izp_sh[base + 1u] = zp.y; + izp_sh[base + 2u] = zp.z; izp_sh[base + 3u] = zp.w; + } + // Group 0 weight sums/scales -> slice 0. + if (gl_LocalInvocationID.x < WG_TILE_N) { + const uint n_idx = tile_n_start + gl_LocalInvocationID.x; + f16vec4 sv = t_weight_scales[n_idx >> 2u]; + wsc_sh[gl_LocalInvocationID.x] = float(sv[n_idx & 3u]); + wsum_sh[gl_LocalInvocationID.x] = t_weight_sums[n_idx]; + } + memoryBarrierShared(); + barrier(); + + // izp/ifs are per-row activation params, constant across K groups — + // broadcast them into coopmats ONCE; the group epilog reuses them every + // group (they depend only on the row block i, not on the group or j). + coopmat + izp_bcast[MMAS_PER_SG_M]; + coopmat + ifs_bcast[MMAS_PER_SG_M]; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + const uint local_m_base = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + coopMatLoad( + izp_bcast[i], izp_sh, + local_m_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutColumnMajor); + coopMatLoad( + ifs_bcast[i], ifs_sh, + local_m_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutColumnMajor); + } + + // dbuf4: prefetch chunk 0 into temp registers, THEN store to slice 0 (no + // barrier here -- the main loop's first iteration barriers before + // reading slice 0). + [[unroll]] for (uint si = 0; si < A_SLOTS_PER_THREAD; ++si) { + const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; + if (slot < A_TOTAL_SLOTS) { + const uint row = slot / K_BLOCKS_PER_CHUNK; + const uint k4 = slot % K_BLOCKS_PER_CHUNK; + temp_A[si] = + t_packed_int8_input[(tile_m_start + row) * nblocks_x_A + k4]; + } + } +#ifdef WEIGHT_INT4 + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { + const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; + const uint block_in_chunk = slot >> 3u; + const uint k4_blk = block_in_chunk / N8_PER_TILE; + const uint n8_blk = (tile_n_start >> 3u) + (block_in_chunk % N8_PER_TILE); +#ifdef WEIGHT_BUFFER + temp_B[si] = t_packed_weight[(n8_blk * nblocks_x_A) + k4_blk]; +#else + temp_B[si] = texelFetch(t_packed_weight, ivec2(k4_blk, n8_blk), 0); +#endif + } +#else + if (b_active) { + const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; +#ifdef WEIGHT_BUFFER + temp_B = t_packed_weight[(b_k4_in_chunk * N4) + block_x_w]; +#else + temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, b_k4_in_chunk), 0); +#endif + } +#endif + { + // store chunk 0 -> slice 0 + [[unroll]] for (uint si = 0; si < A_SLOTS_PER_THREAD; ++si) { + const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; + if (slot < A_TOTAL_SLOTS) { + const uint row = slot / K_BLOCKS_PER_CHUNK; + const uint k4 = slot % K_BLOCKS_PER_CHUNK; + const uint slab_idx = k4 / (MMA_K >> 2u); + const uint k_uint_in_slab = k4 % (MMA_K >> 2u); + Ash_int8 + [slab_idx * A_SLAB_U32 + row * A_STRIDE_U32 + k_uint_in_slab] = + uint(temp_A[si]); + } + } +#ifdef WEIGHT_INT4 + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { + const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; + const uint block_in_chunk = slot >> 3u; + const uint col_in_block = slot & 7u; + const uint k4_in_chunk = block_in_chunk / N8_PER_TILE; + const uint n8_in_tile = block_in_chunk % N8_PER_TILE; + const uint r = col_in_block & 3u; + const uint parity = col_in_block >> 2u; + const int w = temp_B[si][r]; + const int base = int(4u * parity); + const int v0 = (((w >> (base + 0)) & 0xF) - 8) & 0xFF; + const int v1 = (((w >> (base + 8)) & 0xF) - 8) & 0xFF; + const int v2 = (((w >> (base + 16)) & 0xF) - 8) & 0xFF; + const int v3 = (((w >> (base + 24)) & 0xF) - 8) & 0xFF; + const uint n_col = n8_in_tile * 8u + r + parity * 4u; + const uint slab_idx = k4_in_chunk / (MMA_K >> 2u); + const uint k4_in_slab = k4_in_chunk % (MMA_K >> 2u); + Bsh_int8[slab_idx * B_SLAB_U32 + n_col * B_STRIDE_U32 + k4_in_slab] = + uint(v0 | (v1 << 8) | (v2 << 16) | (v3 << 24)); + } +#else + if (b_active) { + const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); + const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); + const uint n_col_base = b_n_uint_col * 4u; + [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { + Bsh_int8[slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab] = + uint(temp_B[n_in_blk]); + } + } +#endif + } + + // ========================================================= + // MAIN LOOP (dbuf4) — nested groups x chunks (kept nested; flattening it + // with a conditional coopmat epilog crashes the Xclipse PAL compiler at + // large spec-resolved trip counts). One barrier per chunk. Chunk + // iteration (global index `chunk`): + // 1. barrier — A/B slice (chunk%2) fully written; on the first chunk + // of group g, wsum/wsc slice (g%2) is too. + // 2. prefetch — chunk+1 (A blocks, B blocks) into temp; when chunk+1 + // starts a new group, also its wsum/wsc element. Skipped + // entirely on the final chunk. + // 3. int8 MMA — on slice (chunk%2) into accum_int32. + // 4. store — temp -> A/B slice ((chunk+1)%2), unpacking the weight; + // on a group boundary, wsum/wsc -> slice ((g+1)%2). + // The group epilog runs unconditionally at the tail of each group. + // ========================================================= + uint chunk = 0; + for (uint group_i = 0; group_i < num_groups; ++group_i) { + for (uint inner = 0; inner < CHUNKS_PER_GROUP; ++inner, ++chunk) { + const bool has_next = chunk + 1u < num_chunks; + const bool group_crossing = has_next && (inner + 1u == CHUNKS_PER_GROUP); + const uint cur_a = (chunk % 2u) * ASH_SLICE_U32; + const uint cur_b = (chunk % 2u) * BSH_SLICE_U32; + const uint nxt_a = ((chunk + 1u) % 2u) * ASH_SLICE_U32; + const uint nxt_b = ((chunk + 1u) % 2u) * BSH_SLICE_U32; + + barrier(); + + // --- 2. prefetch chunk+1 -> temp --- + if (has_next) { + const uint chunkK_nxt = (chunk + 1u) * WG_TILE_K; + [[unroll]] for (uint si = 0; si < A_SLOTS_PER_THREAD; ++si) { + const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; + if (slot < A_TOTAL_SLOTS) { + const uint row = slot / K_BLOCKS_PER_CHUNK; + const uint k4 = slot % K_BLOCKS_PER_CHUNK; + temp_A[si] = t_packed_int8_input + [(tile_m_start + row) * nblocks_x_A + (chunkK_nxt >> 2u) + k4]; + } + } +#ifdef WEIGHT_INT4 + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { + const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; + const uint block_in_chunk = slot >> 3u; + const uint k4_blk = (chunkK_nxt >> 2u) + block_in_chunk / N8_PER_TILE; + const uint n8_blk = (tile_n_start >> 3u) + (block_in_chunk % N8_PER_TILE); +#ifdef WEIGHT_BUFFER + temp_B[si] = t_packed_weight[(n8_blk * nblocks_x_A) + k4_blk]; +#else + temp_B[si] = texelFetch(t_packed_weight, ivec2(k4_blk, n8_blk), 0); +#endif + } + if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { + const uint n_idx = tile_n_start + gl_LocalInvocationID.x; + f16vec4 sv = t_weight_scales[(group_i + 1u) * N4 + (n_idx >> 2u)]; + temp_wsc = float(sv[n_idx & 3u]); + temp_wsum = t_weight_sums[(group_i + 1u) * N + n_idx]; + } +#else + if (b_active) { + const uint block_y_w = (chunkK_nxt >> 2u) + b_k4_in_chunk; + const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; +#ifdef WEIGHT_BUFFER + temp_B = t_packed_weight[(block_y_w * N4) + block_x_w]; +#else + temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, block_y_w), 0); +#endif + } +#endif + } + + // --- 3. int8 MMA on the cur slice --- + [[unroll]] for (uint k = 0; k < NUM_K_SLABS; ++k) { + const uint slab_a_base_u32 = cur_a + k * A_SLAB_U32; + const uint slab_b_base_u32 = cur_b + k * B_SLAB_U32; + + coopmat matA[MMAS_PER_SG_M]; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + const uint row_a = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + coopMatLoad( + matA[i], Ash_int8, + slab_a_base_u32 + row_a * A_STRIDE_U32, + A_STRIDE_U32, + gl_CooperativeMatrixLayoutRowMajor); + } + + coopmat matB; + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint col_b = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopMatLoad( + matB, Bsh_int8, + slab_b_base_u32 + col_b * B_STRIDE_U32, + B_STRIDE_U32, + gl_CooperativeMatrixLayoutColumnMajor); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + accum_int32[i][j] = coopMatMulAdd(matA[i], matB, accum_int32[i][j]); + } + } + } + + // --- 4. store temp (chunk+1) -> nxt slice --- + if (has_next) { + [[unroll]] for (uint si = 0; si < A_SLOTS_PER_THREAD; ++si) { + const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; + if (slot < A_TOTAL_SLOTS) { + const uint row = slot / K_BLOCKS_PER_CHUNK; + const uint k4 = slot % K_BLOCKS_PER_CHUNK; + const uint slab_idx = k4 / (MMA_K >> 2u); + const uint k_uint_in_slab = k4 % (MMA_K >> 2u); + Ash_int8 + [nxt_a + slab_idx * A_SLAB_U32 + row * A_STRIDE_U32 + + k_uint_in_slab] = uint(temp_A[si]); + } + } +#ifdef WEIGHT_INT4 + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { + const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; + const uint block_in_chunk = slot >> 3u; + const uint col_in_block = slot & 7u; + const uint k4_in_chunk = block_in_chunk / N8_PER_TILE; + const uint n8_in_tile = block_in_chunk % N8_PER_TILE; + const uint r = col_in_block & 3u; + const uint parity = col_in_block >> 2u; + const int w = temp_B[si][r]; + const int base = int(4u * parity); + const int v0 = (((w >> (base + 0)) & 0xF) - 8) & 0xFF; + const int v1 = (((w >> (base + 8)) & 0xF) - 8) & 0xFF; + const int v2 = (((w >> (base + 16)) & 0xF) - 8) & 0xFF; + const int v3 = (((w >> (base + 24)) & 0xF) - 8) & 0xFF; + const uint n_col = n8_in_tile * 8u + r + parity * 4u; + const uint slab_idx = k4_in_chunk / (MMA_K >> 2u); + const uint k4_in_slab = k4_in_chunk % (MMA_K >> 2u); + Bsh_int8[nxt_b + slab_idx * B_SLAB_U32 + n_col * B_STRIDE_U32 + k4_in_slab] = + uint(v0 | (v1 << 8) | (v2 << 16) | (v3 << 24)); + } + if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { + const uint wbase_nxt = ((group_i + 1u) % 2u) * WG_TILE_N; + wsum_sh[wbase_nxt + gl_LocalInvocationID.x] = temp_wsum; + wsc_sh[wbase_nxt + gl_LocalInvocationID.x] = temp_wsc; + } +#else + if (b_active) { + const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); + const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); + const uint n_col_base = b_n_uint_col * 4u; + [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { + Bsh_int8[nxt_b + slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab] = + uint(temp_B[n_in_blk]); + } + } +#endif + } + } // chunks + + // --- Group epilog: dequant accum_int32 -> result, reset accum --- + { + const uint wbase = (group_i % 2u) * WG_TILE_N; + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + + coopmat wsum_bcast; + coopMatLoad( + wsum_bcast, wsum_sh, + wbase + local_n_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutRowMajor); + + coopmat wsc_bcast; + coopMatLoad( + wsc_bcast, wsc_sh, + wbase + local_n_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutRowMajor); + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + coopmat adjusted = + accum_int32[i][j] - izp_bcast[i] * wsum_bcast; + coopmat adjusted_fp = + coopmat(adjusted); + coopmat scales_outer = + ifs_bcast[i] * wsc_bcast; + result[i][j] += adjusted_fp * scales_outer; + accum_int32[i][j] = coopmat(0); + } + } + } + } // groups + + // --- Bias (optional) --- +#ifdef HAS_BIAS + if (apply_bias > 0) { + for (uint t = gl_LocalInvocationID.x; t < WG_TILE_N; t += WG_SIZE) { + bias_sh[t] = float(t_bias[tile_n_start + t]); + } + memoryBarrierShared(); + barrier(); + } +#endif + + // --- Store result tile --- + // N for the store address math MUST come from the spec constant, not the + // sizes UBO (see out_N_arg above). +#ifdef IO_TEXTURE + // Epilogue iteration i drains accumulator row-block i from EVERY subgroup + // into Csh_out at once, so the SG_GRID_Y bands it holds are disjoint global + // row ranges; the whole workgroup then imageStores them. lr / MMA_M is the + // writing subgroup's warpInTile.y, so the global row reproduces the buffer + // path's gi exactly. + // + // PORTABILITY NOTE: the barrier() in the loop body keeps this loop rolled + // despite [[unroll]], so result[i][j] IS dynamically indexed. Coopmat arrays + // are opaque per-lane storage and dynamic indexing is exactly the construct + // the Xclipse/AMD-PAL compiler has broken before -- check this first if the + // texture variants miscompile on M51. + const uint CSH_TEXELS_PER_ROW = WG_TILE_N / 4u; + const uint CSH_TEXELS = CSH_ROWS * CSH_TEXELS_PER_ROW; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + // Guards Csh_out against the previous iteration's readers. Inert on i == 0 + // but must stay unconditional to remain workgroup-uniform. + barrier(); + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { +#ifdef HAS_BIAS + if (apply_bias > 0) { + const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopmat bias_tile; + coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); + result[i][j] += bias_tile; + } +#endif + coopmat out_tile = + coopmat(result[i][j]); + coopMatStore( + out_tile, Csh_out, + warpInTile.y * MMA_M * WG_TILE_N + + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j), + WG_TILE_N, + gl_CooperativeMatrixLayoutRowMajor); + } + memoryBarrierShared(); + barrier(); + + for (uint t = gl_LocalInvocationID.x; t < CSH_TEXELS; t += WG_SIZE) { + const uint lr = t / CSH_TEXELS_PER_ROW; + const uint lc4 = t % CSH_TEXELS_PER_ROW; + const uint m = + tile_m_start + (lr / MMA_M) * SG_TILE_M + i * MMA_M + (lr % MMA_M); + const uint base = lr * WG_TILE_N + lc4 * 4u; + imageStore( + t_output, + ivec3(tile_n_start / 4u + lc4, m, 0), + vec4( + float(Csh_out[base]), + float(Csh_out[base + 1u]), + float(Csh_out[base + 2u]), + float(Csh_out[base + 3u]))); + } + } +#else + const uint N_out = uint(out_N_arg); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint gi = tile_m_start + MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + const uint gj = tile_n_start + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + +#ifdef HAS_BIAS + if (apply_bias > 0) { + const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopmat bias_tile; + coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); + result[i][j] += bias_tile; + } +#endif + + coopmat out_tile = + coopmat(result[i][j]); + coopMatStore( + out_tile, t_output, + gi * N_out + gj, N_out, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#endif // IO_TEXTURE +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trm.yaml b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trm.yaml new file mode 100644 index 00000000000..28f9bdf03e5 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trm.yaml @@ -0,0 +1,188 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# "-tr" (coopmat-staged A) variants of the int8 dq8ca_q4gsw coopmat kernel, +# ported from shmem_double_buf4-tr.comp (vk_cooperative_matrix_perf, +# gemm-ubm). linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trm.glsl is a fork of +# linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4.glsl with ONLY the A-side global -> +# LDS staging swapped from a per-thread ivec4 copy to +# coopMatLoad(global) -> coopmat<> -> coopMatStore(LDS). B staging is +# byte-identical (int4 nibble unpack cannot be coopmat-staged) and B was +# already column-major in LDS, so it is not a delta. +# +# These variants require t_packed_int8_input in the ROW-MAJOR +# kPackedInt8_4W layout (produced by quantize_and_pack_4w_with_group_sums), +# NOT the stock 4h4w block layout. QuantizedLinear.cpp switches the tensor +# layout and the packer node together with the variant token, and only when +# the coopmat gate passes -- see dq8ca_wants_rowmajor_int8_input(). +# +# Selected via +# ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4trm_txkgs. +# NOT the shipped default: unvalidated until it passes repeated +# test_llama_microbench --correctness-only runs (a single pass is not proof -- +# see dq8ca_coopmat_variant()'s comment). +# +# Tile-geometry preconditions beyond dbuf4's: WG_TILE_M % MMA_M == 0 and +# WG_TILE_K % MMA_K == 0 (both hold for every entry below). The seed set is +# deliberately small -- widen it from a sweep once correctness is established. + +linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trm: + parameter_names_with_default_values: + PRECISION: highp + HAS_BIAS: false + IO_STORAGE: buffer + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + MMA_M: 16 + MMA_N: 16 + MMA_K: 16 + WG_TILE_M: 64 + WG_TILE_N: 32 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + shader_variants: + # production 8da4w tile (dbuf2/dbuf4 default) -- the A/B anchor + # A tiles/chunk = 8, subgroups = 2 -> A_TILES_PER_SG = 4 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trm_t64x32k32g12s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 32 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trm_t64x32k32g12s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 32 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trm_t64x32k32g12s64_buffer_buffer_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: buffer + WG_TILE_M: 64 + WG_TILE_N: 32 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + # q4gsw production geometry; 8 A-tiles over 4 subgroups + # A tiles/chunk = 8, subgroups = 4 -> A_TILES_PER_SG = 2 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trm_t128x128k16g22s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trm_t128x128k16g22s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + # 1 A-tile per subgroup (A_TILES_PER_SG == 1, no leftover slot) + # A tiles/chunk = 8, subgroups = 8 -> A_TILES_PER_SG = 1 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trm_t64x64k32g24s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trm_t64x64k32g24s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + # 16 A-tiles over 8 subgroups + # A tiles/chunk = 16, subgroups = 8 -> A_TILES_PER_SG = 2 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trm_t128x64k32g24s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trm_t128x64k32g24s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 2 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + # Reference-matching tile: shmem_double_buf4-tr.comp's own geometry + # (TILE 128x128x16, WORKGROUP_WIDTH_IN_SUBGROUPS=4 x DBUF4_GRID_HEIGHT=2, + # wave32) => 8 subgroups / 256 threads, MMAS_PER_SG 4x2. Added so our + # kernel can be A/B'd against the teammate's PAL capture at the same + # tile and the same 2048x1024x4096 shape (8b wk_wv prefill). + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trm_t128x128k16g42s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trm_t128x128k16g42s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + # G2/Pavan's documented INT8 uBM config: TILE 128x128x64 with a 4x2 + # subgroup grid at wave32 (workgroupSize=256). Added to compare our kernel + # against their reference dbuf4 number (~1973us) at the same geometry and + # the same shape (M2048 K4096 N1024 == 8b wk_wv prefill). + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trm_t128x128k64g42s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trm_t128x128k64g42s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp.glsl b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp.glsl new file mode 100644 index 00000000000..65615a2f1f5 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp.glsl @@ -0,0 +1,687 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +/* + * "zp-hoisted" variant: identical to dbuf4 except that the activation + * zero-point correction and the per-row activation scale are applied ONCE + * after the group loop instead of once per quantization group. + * + * The per-group epilogue term factors exactly: + * + * out[m][n] = ifs[m] * SUM_g wsc[g][n] * ( acc[g] - izp[m]*wsum[g][n] ) + * = ifs[m] * [ SUM_g wsc[g][n]*acc[g] - izp[m]*SUM_g wsc[g][n]*wsum[g][n] ] + * \__ weight-side only __/ + * + * `ifs` is per-row and group-independent so it factors out entirely, and the + * zero-point term separates into a per-row scalar times a per-output-channel + * weight-side sum. That sum depends on no activation data, so it is + * accumulated once into `wcorr_sh` in the prologue rather than being rebuilt + * per group. + * + * Consequences vs dbuf4, per accumulator tile per group: + * - gone: izp*wsum multiply and the subtract (48 v_sub* in the loop body) + * - gone: ifs*wsc multiply (part of 48 dequant-fp) + * - gone: the wsum_sh shared array and its ping-pong + * - gone: izp_bcast / ifs_bcast live across the loop (register pressure) + * - kept: result += float(acc) * wsc + * + * Exact in exact arithmetic, but NOT bit-exact in fp32 -- the summation order + * changes -- so it is gated on the correctness matrix like any other change. + * + * No new binding and no export-format change: the weight-side sum is derived + * in the prologue from t_weight_scales and t_weight_sums, both already bound. + * + * Selected via + * ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4zp_txkgs<32|64>. + * + * Original dbuf4 header follows. + * + * TILE/SUBGROUP-SWEEP variant of the int8 dq8ca_q4gsw coopmat shader's dbuf4 + * ("store-first-for-next", the ORIGINAL loop structure before specs/025 User + * Story 1 picked dbuf2) loop structure (specs/041-dbuf4-tile-sweep). Forked + * from linear_dq8ca_q4gsw_coopmat_tsweep.glsl (which carries dbuf2's loop, + * the production winner) -- everything except the PROLOGUE/MAIN LOOP block + * is identical: bindings, spec-constants, tile-geometry templating, LDS + * layout (ColumnMajor B + skew), int8 WMMA thread maps, group epilog, + * bias/store epilogue. Only the loop structure is swapped to dbuf4, + * recovered from git commit 8d0f23ee78's + * linear_dq8ca_q4gsw_coopmat_dbuf4.glsl (see specs/041/reference/) -- the + * byte-identical pre-swap copy of what is now linear_dq8ca_qw_coopmat.glsl. + * + * The nested `groups x chunks` loop and unconditional group epilog are kept + * exactly as in dbuf2 -- flattening them crashes the Xclipse PAL compiler at + * large spec-resolved trip counts (see dbuf2's own header). Only the + * store/barrier/prefetch ORDER within each chunk iteration is inverted: + * + * dbuf2 (this file's base): store(temp, already prefetched -> cur slice) + * -> barrier -> MMA(cur) -> prefetch(next -> temp) [store owns the + * CURRENT chunk, at the iteration's start] + * dbuf4 (this file): barrier -> prefetch(next -> temp) -> MMA(cur) -> + * store(temp -> next slice) [store owns the NEXT chunk, at the + * iteration's end -- the mirror image] + * + * The group wsum/wsc ping-pong is inverted the same way: dbuf4 stores the + * next group's values (prefetched during the crossing chunk) at the TAIL of + * that chunk, instead of dbuf2's HEAD-of-new-group placement. + * + * Selected at dispatch via + * ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4_txkgs<32|64> + * (QuantizedLinear.cpp), additive to the existing tsweep_t... (dbuf2) + * namespace. + * + * KHR Cooperative Matrix variant of the dynamically-quantized-activation + * linear tiled shader (WEIGHT_NBITS=4): + * 4 -> linear_dq8ca_q4gsw_coopmat INT4 group-symmetric weight + * + * Performs: out[M,N] = dequant(int8_act) * dequant(int_w) (+ bias) + * via coopmat x coopmat -> coopmat on the matrix unit. + * + * Hard preconditions: + * M % WG_TILE_M == 0, N % WG_TILE_N == 0, K % WG_TILE_K == 0, + * INT4: group_size % WG_TILE_K == 0, + * device exposes coopmatx-> at 16x16x16. + */ + +#version 450 core + +#extension GL_KHR_cooperative_matrix : require +#extension GL_KHR_memory_scope_semantics : require +#extension GL_KHR_shader_subgroup_basic : enable +#extension GL_EXT_shader_explicit_arithmetic_types : require +#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_EXT_control_flow_attributes : enable + +#define PRECISION ${PRECISION} + +$if WEIGHT_NBITS == 4: + #define WEIGHT_INT4 + +$if HAS_BIAS: + #define HAS_BIAS + +$if WEIGHT_STORAGE == "buffer": + #define WEIGHT_BUFFER + +$if IO_STORAGE == "texture3d": + #define IO_TEXTURE + +layout(std430) buffer; + +#include "common.glslh" + +// Bindings — match add_linear_dqa_qw_node arg order: +// output(0), fp_input(1), packed_int8_input(2), int_input_sums(3 - unused), +// input_scales(4), input_zps(5), packed_weight(6), weight_sums(7), +// weight_scales(8), bias(9). +${layout_declare_tensor(B, "w", "t_output", "half", IO_STORAGE, is_scalar_array=True)} +// t_input is unread here -- the activations arrive already quantized in +// t_packed_int8_input -- but stays declared so the binding layout matches the +// dispatch site. It tracks IO_STORAGE so the two IO tensors stay consistent. +${layout_declare_tensor(B, "r", "t_input", "half", IO_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_packed_int8_input", "int", "buffer", is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_int8_input_sums", "int", "buffer", is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_int8_input_scales", "half", "texture3d")} +${layout_declare_tensor(B, "r", "t_int8_input_zps", "int8", "texture3d")} +${layout_declare_tensor(B, "r", "t_packed_weight", "int", WEIGHT_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_weight_sums", "int", "buffer", is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_weight_scales", "half", "buffer", is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_bias", "half", "buffer", is_scalar_array=True)} + +${layout_declare_ubo(B, "ivec4", "output_sizes")} +${layout_declare_ubo(B, "ivec4", "input_sizes")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +${layout_declare_spec_const(C, "int", "apply_bias", "0")} +// INT4 only; inert (0) for INT8 so the dispatcher's spec list lines up. +${layout_declare_spec_const(C, "int", "K4_per_group", "0")} +${layout_declare_spec_const(C, "int", "num_groups_arg", "0")} +${layout_declare_spec_const(C, "int", "out_N_arg", "0")} + +// Tile geometry +const uint MMA_M = ${MMA_M}; +const uint MMA_N = ${MMA_N}; +const uint MMA_K = ${MMA_K}; + +const uint WG_TILE_M = ${WG_TILE_M}; +const uint WG_TILE_N = ${WG_TILE_N}; +const uint WG_TILE_K = ${WG_TILE_K}; + +const uint SG_GRID_X = ${SG_GRID_X}; +const uint SG_GRID_Y = ${SG_GRID_Y}; +const uint SUBGROUP_SIZE = ${SUBGROUP_SIZE}; +const uint NUM_SUBGROUPS = SG_GRID_X * SG_GRID_Y; +const uint WG_SIZE = NUM_SUBGROUPS * SUBGROUP_SIZE; + +const uint SG_TILE_M = WG_TILE_M / SG_GRID_Y; +const uint SG_TILE_N = WG_TILE_N / SG_GRID_X; +const uint MMAS_PER_SG_M = SG_TILE_M / MMA_M; +const uint MMAS_PER_SG_N = SG_TILE_N / MMA_N; + +const uint A_SLAB_INT8 = WG_TILE_M * MMA_K; +const uint B_USEFUL_U32 = MMA_K / 4u; +const uint B_STRIDE_U32 = B_USEFUL_U32 + 1u; // +1 skew +const uint B_SLAB_U32 = WG_TILE_N * B_STRIDE_U32; +const uint NUM_K_SLABS = WG_TILE_K / MMA_K; + +const uint A_SLAB_U32 = A_SLAB_INT8 / 4u; +const uint A_STRIDE_U32 = MMA_K / 4u; + +// One ping-pong slice covers all K-slabs of one chunk. +const uint ASH_SLICE_U32 = NUM_K_SLABS * A_SLAB_U32; +const uint BSH_SLICE_U32 = NUM_K_SLABS * B_SLAB_U32; + +// Double-buffered MMA operand staging. +shared uint Ash_int8[2u * ASH_SLICE_U32]; +shared uint Bsh_int8[2u * BSH_SLICE_U32]; + +// Per-WG-tile-row activation params (loaded ONCE at WG start; constant +// across groups). +shared int izp_sh[WG_TILE_M]; // int32 (cast from int8 source) for broadcast +shared float ifs_sh[WG_TILE_M]; // float32 (cast from fp16 source) for broadcast + +// Per-(group, output-channel) weight params, ping-ponged by group parity. +// (For per-channel INT8 only slice 0 is ever used.) +shared float wsc_sh[2u * WG_TILE_N]; +// SUM_g wsc[g][n]*wsum[g][n] per output channel -- weight-side only, so it is +// accumulated once in the prologue. Replaces dbuf4's ping-ponged wsum_sh +// (which was 2*WG_TILE_N ints), so this is a net LDS saving. +shared float wcorr_sh[WG_TILE_N]; + +#ifdef HAS_BIAS +shared float bias_sh[WG_TILE_N]; +#endif + +#ifdef IO_TEXTURE +// Result staging for the imageStore epilogue, mirroring the fp16 kernel: +// SG_GRID_Y bands of MMA_M rows, each WG_TILE_N wide, row-major. A full +// WG_TILE_M x WG_TILE_N buffer would cost SG_GRID_Y/MMAS_PER_SG_M x more LDS +// and wreck occupancy. float16_t-typed because coopMatStore needs it. +const uint CSH_ROWS = SG_GRID_Y * MMA_M; +shared float16_t Csh_out[CSH_ROWS * WG_TILE_N]; +#endif + +// Running fp32 accumulator (across all groups). +coopmat + result[MMAS_PER_SG_M][MMAS_PER_SG_N]; + +// Per-group int32 MMA accumulator. +coopmat + accum_int32[MMAS_PER_SG_M][MMAS_PER_SG_N]; + +void main() { + const uvec2 tileID = uvec2(gl_WorkGroupID.xy); + const uvec2 warpInTile = uvec2( + gl_SubgroupID % SG_GRID_X, + gl_SubgroupID / SG_GRID_X); + + const uint K = uint(input_sizes.x); + const uint N = uint(output_sizes.x); + const uint N4 = (N + 3u) / 4u; + const uint nblocks_x_A = (K + 3u) >> 2u; + +#ifdef WEIGHT_INT4 + const uint num_groups = uint(num_groups_arg); + const uint CHUNKS_PER_GROUP = uint(K4_per_group) * 4u / WG_TILE_K; +#else + // Per-channel: a single quant "group" spanning all of K. The nested + // groups x chunks loop below collapses to a flat chunk loop, the wsum/wsc + // ping-pong never crosses a boundary, and the epilog runs exactly once. + const uint num_groups = 1u; + const uint CHUNKS_PER_GROUP = uint(num_groups_arg); +#endif + const uint num_chunks = num_groups * CHUNKS_PER_GROUP; + + const uint tile_m_start = WG_TILE_M * tileID.y; + const uint tile_n_start = WG_TILE_N * tileID.x; + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + result[i][j] = coopmat(0.0); + accum_int32[i][j] = coopmat(0); + } + } + + // --- A staging thread map: one (m4, k4) ivec4 block per active thread --- + const uint K_BLOCKS_PER_CHUNK = WG_TILE_K >> 2u; + const uint A_ACTIVE_THREADS = (WG_TILE_M >> 2u) * K_BLOCKS_PER_CHUNK; + const uint a_m_block = gl_LocalInvocationID.x / K_BLOCKS_PER_CHUNK; + const uint a_k_block = gl_LocalInvocationID.x % K_BLOCKS_PER_CHUNK; + const bool a_active = gl_LocalInvocationID.x < A_ACTIVE_THREADS; + +#ifdef WEIGHT_INT4 + // --- B staging thread map: (block, col) slots; each slot extracts one + // ColumnMajor LDS uint (4 K-contiguous sign-extended int8) --- + const uint B_TOTAL_SLOTS = K_BLOCKS_PER_CHUNK * WG_TILE_N; + const uint B_SLOTS_PER_THREAD = B_TOTAL_SLOTS / WG_SIZE; + const uint N8_PER_TILE = WG_TILE_N >> 3u; +#else + // --- B staging thread map: one (k4, n4) ivec4 block per active thread --- + const uint B_FETCH_SLOTS = K_BLOCKS_PER_CHUNK * (WG_TILE_N >> 2u); + const uint N4_PER_TILE = WG_TILE_N >> 2u; + const uint b_k4_in_chunk = gl_LocalInvocationID.x / N4_PER_TILE; + const uint b_n_uint_col = gl_LocalInvocationID.x % N4_PER_TILE; + const bool b_active = gl_LocalInvocationID.x < B_FETCH_SLOTS; +#endif + + // Prefetch temp registers. + ivec4 temp_A; +#ifdef WEIGHT_INT4 + ivec4 temp_B[B_SLOTS_PER_THREAD]; + float temp_wsc; +#else + ivec4 temp_B; +#endif + + // ========================================================= + // PROLOGUE + // ========================================================= + if (gl_LocalInvocationID.x < (WG_TILE_M >> 2u)) { + const uint m4 = (tile_m_start >> 2u) + gl_LocalInvocationID.x; + const vec4 sc = vec4(texelFetch(t_int8_input_scales, ivec3(m4, 0, 0), 0)); + const ivec4 zp = texelFetch(t_int8_input_zps, ivec3(m4, 0, 0), 0); + const uint base = gl_LocalInvocationID.x * 4u; + ifs_sh[base + 0u] = sc.x; ifs_sh[base + 1u] = sc.y; + ifs_sh[base + 2u] = sc.z; ifs_sh[base + 3u] = sc.w; + izp_sh[base + 0u] = zp.x; izp_sh[base + 1u] = zp.y; + izp_sh[base + 2u] = zp.z; izp_sh[base + 3u] = zp.w; + } + // Group 0 weight scales -> slice 0, and the hoisted weight-side correction + // SUM_g wsc[g][n]*wsum[g][n] accumulated across ALL groups. The loop is + // prologue-only (the prologue is ~1.4% of the dynamic instruction stream), + // and it replaces per-group wsum work inside the loop body. + if (gl_LocalInvocationID.x < WG_TILE_N) { + const uint n_idx = tile_n_start + gl_LocalInvocationID.x; + f16vec4 sv0 = t_weight_scales[n_idx >> 2u]; + wsc_sh[gl_LocalInvocationID.x] = float(sv0[n_idx & 3u]); + + float corr = 0.0; + for (uint g = 0; g < num_groups; ++g) { + f16vec4 sv = t_weight_scales[g * N4 + (n_idx >> 2u)]; + corr += float(sv[n_idx & 3u]) * float(t_weight_sums[g * N + n_idx]); + } + wcorr_sh[gl_LocalInvocationID.x] = corr; + } + memoryBarrierShared(); + barrier(); + + // NOTE: dbuf4 builds izp_bcast/ifs_bcast here and keeps them live across the + // whole group loop. This variant needs them only AFTER the loop, so they are + // loaded there instead -- that is the register-pressure saving. + + // dbuf4: prefetch chunk 0 into temp registers, THEN store to slice 0 (no + // barrier here -- the main loop's first iteration barriers before + // reading slice 0). + if (a_active) { + const uint m4_global = (tile_m_start >> 2u) + a_m_block; + temp_A = t_packed_int8_input[m4_global * nblocks_x_A + a_k_block]; + } +#ifdef WEIGHT_INT4 + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { + const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; + const uint block_in_chunk = slot >> 3u; + const uint k4_blk = block_in_chunk / N8_PER_TILE; + const uint n8_blk = (tile_n_start >> 3u) + (block_in_chunk % N8_PER_TILE); +#ifdef WEIGHT_BUFFER + temp_B[si] = t_packed_weight[(n8_blk * nblocks_x_A) + k4_blk]; +#else + temp_B[si] = texelFetch(t_packed_weight, ivec2(k4_blk, n8_blk), 0); +#endif + } +#else + if (b_active) { + const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; +#ifdef WEIGHT_BUFFER + temp_B = t_packed_weight[(b_k4_in_chunk * N4) + block_x_w]; +#else + temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, b_k4_in_chunk), 0); +#endif + } +#endif + { + // store chunk 0 -> slice 0 + if (a_active) { + const uint slab_idx = a_k_block / (MMA_K >> 2u); + const uint k_uint_in_slab = a_k_block % (MMA_K >> 2u); + const uint base_row = a_m_block * 4u; + [[unroll]] for (uint m4i = 0; m4i < 4u; ++m4i) { + Ash_int8[slab_idx * A_SLAB_U32 + (base_row + m4i) * A_STRIDE_U32 + k_uint_in_slab] = + uint(temp_A[m4i]); + } + } +#ifdef WEIGHT_INT4 + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { + const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; + const uint block_in_chunk = slot >> 3u; + const uint col_in_block = slot & 7u; + const uint k4_in_chunk = block_in_chunk / N8_PER_TILE; + const uint n8_in_tile = block_in_chunk % N8_PER_TILE; + const uint r = col_in_block & 3u; + const uint parity = col_in_block >> 2u; + const int w = temp_B[si][r]; + const int base = int(4u * parity); + const int v0 = (((w >> (base + 0)) & 0xF) - 8) & 0xFF; + const int v1 = (((w >> (base + 8)) & 0xF) - 8) & 0xFF; + const int v2 = (((w >> (base + 16)) & 0xF) - 8) & 0xFF; + const int v3 = (((w >> (base + 24)) & 0xF) - 8) & 0xFF; + const uint n_col = n8_in_tile * 8u + r + parity * 4u; + const uint slab_idx = k4_in_chunk / (MMA_K >> 2u); + const uint k4_in_slab = k4_in_chunk % (MMA_K >> 2u); + Bsh_int8[slab_idx * B_SLAB_U32 + n_col * B_STRIDE_U32 + k4_in_slab] = + uint(v0 | (v1 << 8) | (v2 << 16) | (v3 << 24)); + } +#else + if (b_active) { + const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); + const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); + const uint n_col_base = b_n_uint_col * 4u; + [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { + Bsh_int8[slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab] = + uint(temp_B[n_in_blk]); + } + } +#endif + } + + // ========================================================= + // MAIN LOOP (dbuf4) — nested groups x chunks (kept nested; flattening it + // with a conditional coopmat epilog crashes the Xclipse PAL compiler at + // large spec-resolved trip counts). One barrier per chunk. Chunk + // iteration (global index `chunk`): + // 1. barrier — A/B slice (chunk%2) fully written; on the first chunk + // of group g, wsum/wsc slice (g%2) is too. + // 2. prefetch — chunk+1 (A blocks, B blocks) into temp; when chunk+1 + // starts a new group, also its wsum/wsc element. Skipped + // entirely on the final chunk. + // 3. int8 MMA — on slice (chunk%2) into accum_int32. + // 4. store — temp -> A/B slice ((chunk+1)%2), unpacking the weight; + // on a group boundary, wsum/wsc -> slice ((g+1)%2). + // The group epilog runs unconditionally at the tail of each group. + // ========================================================= + uint chunk = 0; + for (uint group_i = 0; group_i < num_groups; ++group_i) { + for (uint inner = 0; inner < CHUNKS_PER_GROUP; ++inner, ++chunk) { + const bool has_next = chunk + 1u < num_chunks; + const bool group_crossing = has_next && (inner + 1u == CHUNKS_PER_GROUP); + const uint cur_a = (chunk % 2u) * ASH_SLICE_U32; + const uint cur_b = (chunk % 2u) * BSH_SLICE_U32; + const uint nxt_a = ((chunk + 1u) % 2u) * ASH_SLICE_U32; + const uint nxt_b = ((chunk + 1u) % 2u) * BSH_SLICE_U32; + + barrier(); + + // --- 2. prefetch chunk+1 -> temp --- + if (has_next) { + const uint chunkK_nxt = (chunk + 1u) * WG_TILE_K; + if (a_active) { + const uint m4_global = (tile_m_start >> 2u) + a_m_block; + const uint k4_global = (chunkK_nxt >> 2u) + a_k_block; + temp_A = t_packed_int8_input[m4_global * nblocks_x_A + k4_global]; + } +#ifdef WEIGHT_INT4 + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { + const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; + const uint block_in_chunk = slot >> 3u; + const uint k4_blk = (chunkK_nxt >> 2u) + block_in_chunk / N8_PER_TILE; + const uint n8_blk = (tile_n_start >> 3u) + (block_in_chunk % N8_PER_TILE); +#ifdef WEIGHT_BUFFER + temp_B[si] = t_packed_weight[(n8_blk * nblocks_x_A) + k4_blk]; +#else + temp_B[si] = texelFetch(t_packed_weight, ivec2(k4_blk, n8_blk), 0); +#endif + } + if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { + const uint n_idx = tile_n_start + gl_LocalInvocationID.x; + f16vec4 sv = t_weight_scales[(group_i + 1u) * N4 + (n_idx >> 2u)]; + temp_wsc = float(sv[n_idx & 3u]); + } +#else + if (b_active) { + const uint block_y_w = (chunkK_nxt >> 2u) + b_k4_in_chunk; + const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; +#ifdef WEIGHT_BUFFER + temp_B = t_packed_weight[(block_y_w * N4) + block_x_w]; +#else + temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, block_y_w), 0); +#endif + } +#endif + } + + // --- 3. int8 MMA on the cur slice --- + [[unroll]] for (uint k = 0; k < NUM_K_SLABS; ++k) { + const uint slab_a_base_u32 = cur_a + k * A_SLAB_U32; + const uint slab_b_base_u32 = cur_b + k * B_SLAB_U32; + + coopmat matA[MMAS_PER_SG_M]; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + const uint row_a = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + coopMatLoad( + matA[i], Ash_int8, + slab_a_base_u32 + row_a * A_STRIDE_U32, + A_STRIDE_U32, + gl_CooperativeMatrixLayoutRowMajor); + } + + coopmat matB; + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint col_b = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopMatLoad( + matB, Bsh_int8, + slab_b_base_u32 + col_b * B_STRIDE_U32, + B_STRIDE_U32, + gl_CooperativeMatrixLayoutColumnMajor); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + accum_int32[i][j] = coopMatMulAdd(matA[i], matB, accum_int32[i][j]); + } + } + } + + // --- 4. store temp (chunk+1) -> nxt slice --- + if (has_next) { + if (a_active) { + const uint slab_idx = a_k_block / (MMA_K >> 2u); + const uint k_uint_in_slab = a_k_block % (MMA_K >> 2u); + const uint base_row = a_m_block * 4u; + [[unroll]] for (uint m4i = 0; m4i < 4u; ++m4i) { + Ash_int8[nxt_a + slab_idx * A_SLAB_U32 + (base_row + m4i) * A_STRIDE_U32 + k_uint_in_slab] = + uint(temp_A[m4i]); + } + } +#ifdef WEIGHT_INT4 + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { + const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; + const uint block_in_chunk = slot >> 3u; + const uint col_in_block = slot & 7u; + const uint k4_in_chunk = block_in_chunk / N8_PER_TILE; + const uint n8_in_tile = block_in_chunk % N8_PER_TILE; + const uint r = col_in_block & 3u; + const uint parity = col_in_block >> 2u; + const int w = temp_B[si][r]; + const int base = int(4u * parity); + const int v0 = (((w >> (base + 0)) & 0xF) - 8) & 0xFF; + const int v1 = (((w >> (base + 8)) & 0xF) - 8) & 0xFF; + const int v2 = (((w >> (base + 16)) & 0xF) - 8) & 0xFF; + const int v3 = (((w >> (base + 24)) & 0xF) - 8) & 0xFF; + const uint n_col = n8_in_tile * 8u + r + parity * 4u; + const uint slab_idx = k4_in_chunk / (MMA_K >> 2u); + const uint k4_in_slab = k4_in_chunk % (MMA_K >> 2u); + Bsh_int8[nxt_b + slab_idx * B_SLAB_U32 + n_col * B_STRIDE_U32 + k4_in_slab] = + uint(v0 | (v1 << 8) | (v2 << 16) | (v3 << 24)); + } + if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { + const uint wbase_nxt = ((group_i + 1u) % 2u) * WG_TILE_N; + wsc_sh[wbase_nxt + gl_LocalInvocationID.x] = temp_wsc; + } +#else + if (b_active) { + const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); + const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); + const uint n_col_base = b_n_uint_col * 4u; + [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { + Bsh_int8[nxt_b + slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab] = + uint(temp_B[n_in_blk]); + } + } +#endif + } + } // chunks + + // --- Group epilog: scale-only accumulate, reset accum --- + // Just result += float(acc) * wsc. The zero-point subtract and the ifs + // multiply are hoisted out of the group loop (applied once below). + { + const uint wbase = (group_i % 2u) * WG_TILE_N; + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + + coopmat wsc_bcast; + coopMatLoad( + wsc_bcast, wsc_sh, + wbase + local_n_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutRowMajor); + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + result[i][j] += + coopmat( + accum_int32[i][j]) * wsc_bcast; + accum_int32[i][j] = coopmat(0); + } + } + } + } // groups + + // --- Hoisted correction, applied ONCE: --------------------------------- + // result = ifs * ( result - izp * SUM_g wsc_g*wsum_g ) + // izp/ifs are loaded here rather than before the group loop so they are not + // live across it. + { + coopmat + izpf_bcast[MMAS_PER_SG_M]; + coopmat + ifs_bcast[MMAS_PER_SG_M]; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + const uint local_m_base = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + coopmat izp_i; + coopMatLoad( + izp_i, izp_sh, local_m_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutColumnMajor); + izpf_bcast[i] = + coopmat(izp_i); + coopMatLoad( + ifs_bcast[i], ifs_sh, local_m_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutColumnMajor); + } + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopmat wcorr_bcast; + coopMatLoad( + wcorr_bcast, wcorr_sh, local_n_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutRowMajor); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + result[i][j] = ifs_bcast[i] * (result[i][j] - izpf_bcast[i] * wcorr_bcast); + } + } + } + + // --- Bias (optional) --- +#ifdef HAS_BIAS + if (apply_bias > 0) { + for (uint t = gl_LocalInvocationID.x; t < WG_TILE_N; t += WG_SIZE) { + bias_sh[t] = float(t_bias[tile_n_start + t]); + } + memoryBarrierShared(); + barrier(); + } +#endif + + // --- Store result tile --- + // N for the store address math MUST come from the spec constant, not the + // sizes UBO (see out_N_arg above). +#ifdef IO_TEXTURE + // Epilogue iteration i drains accumulator row-block i from EVERY subgroup + // into Csh_out at once, so the SG_GRID_Y bands it holds are disjoint global + // row ranges; the whole workgroup then imageStores them. lr / MMA_M is the + // writing subgroup's warpInTile.y, so the global row reproduces the buffer + // path's gi exactly. + // + // PORTABILITY NOTE: the barrier() in the loop body keeps this loop rolled + // despite [[unroll]], so result[i][j] IS dynamically indexed. Coopmat arrays + // are opaque per-lane storage and dynamic indexing is exactly the construct + // the Xclipse/AMD-PAL compiler has broken before -- check this first if the + // texture variants miscompile on M51. + const uint CSH_TEXELS_PER_ROW = WG_TILE_N / 4u; + const uint CSH_TEXELS = CSH_ROWS * CSH_TEXELS_PER_ROW; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + // Guards Csh_out against the previous iteration's readers. Inert on i == 0 + // but must stay unconditional to remain workgroup-uniform. + barrier(); + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { +#ifdef HAS_BIAS + if (apply_bias > 0) { + const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopmat bias_tile; + coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); + result[i][j] += bias_tile; + } +#endif + coopmat out_tile = + coopmat(result[i][j]); + coopMatStore( + out_tile, Csh_out, + warpInTile.y * MMA_M * WG_TILE_N + + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j), + WG_TILE_N, + gl_CooperativeMatrixLayoutRowMajor); + } + memoryBarrierShared(); + barrier(); + + for (uint t = gl_LocalInvocationID.x; t < CSH_TEXELS; t += WG_SIZE) { + const uint lr = t / CSH_TEXELS_PER_ROW; + const uint lc4 = t % CSH_TEXELS_PER_ROW; + const uint m = + tile_m_start + (lr / MMA_M) * SG_TILE_M + i * MMA_M + (lr % MMA_M); + const uint base = lr * WG_TILE_N + lc4 * 4u; + imageStore( + t_output, + ivec3(tile_n_start / 4u + lc4, m, 0), + vec4( + float(Csh_out[base]), + float(Csh_out[base + 1u]), + float(Csh_out[base + 2u]), + float(Csh_out[base + 3u]))); + } + } +#else + const uint N_out = uint(out_N_arg); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint gi = tile_m_start + MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + const uint gj = tile_n_start + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + +#ifdef HAS_BIAS + if (apply_bias > 0) { + const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopmat bias_tile; + coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); + result[i][j] += bias_tile; + } +#endif + + coopmat out_tile = + coopmat(result[i][j]); + coopMatStore( + out_tile, t_output, + gi * N_out + gj, N_out, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#endif // IO_TEXTURE +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp.yaml b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp.yaml new file mode 100644 index 00000000000..40e7a12f29d --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp.yaml @@ -0,0 +1,187 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# "zp-hoisted" variants of the int8 dq8ca_q4gsw coopmat kernel: identical to +# tsweep_dbuf4 except the activation zero-point correction and the per-row +# activation scale are applied once after the group loop instead of once per +# quantization group. No new binding, no export-format change -- the weight-side +# correction sum is derived in the prologue from tensors already bound. +# +# Selected via +# ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4zp_txkgs. +# NOT the shipped default until it passes repeated correctness runs. + +linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp: + parameter_names_with_default_values: + PRECISION: highp + HAS_BIAS: false + IO_STORAGE: buffer + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + MMA_M: 16 + MMA_N: 16 + MMA_K: 16 + WG_TILE_M: 64 + WG_TILE_N: 32 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + shader_variants: + # production tile -- the A/B anchor against the 1.3 baseline + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp_t64x32k32g12s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 32 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp_t64x32k32g12s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 32 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp_t64x32k32g12s64_buffer_buffer_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: buffer + WG_TILE_M: 64 + WG_TILE_N: 32 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + # NOTE: t128x128k64g42s32 was seeded here and then REMOVED. Its A-staging + # thread map cannot cover its blocks: A_ACTIVE_THREADS = (128/4)*(64/4) = + # 512 but WG_SIZE = 4*2*32 = 256, so half of A is never staged and the + # kernel is numerically wrong (measured 12/12 correctness failures on the + # dbuf4 equivalent). Use g44s32 for that tile instead -- 4*4*32 = 512 = A. + # deep K + big tile; the reference's winning geometry, A map valid at 512=512 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp_t128x128k64g44s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 4 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp_t128x128k64g44s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 4 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + # task 5.1's explicit ask + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp_t128x128k32g42s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp_t128x128k32g42s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + # same tile, more subgroups -> MMAS 2x2 instead of 4x2 (less accumulator pressure) + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp_t128x128k32g44s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp_t128x128k32g44s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + # moderate step up from the shipped 64x32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp_t128x64k32g42s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp_t128x64k32g42s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + # small step up + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp_t64x64k32g22s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp_t64x64k32g22s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + # big tile, shallow K -- isolates tile area from K depth + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp_t128x128k16g42s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp_t128x128k16g42s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpb.glsl b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpb.glsl new file mode 100644 index 00000000000..1b90a3bb647 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpb.glsl @@ -0,0 +1,714 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +/* + * "zp-hoisted" variant: identical to dbuf4 except that the activation + * zero-point correction and the per-row activation scale are applied ONCE + * after the group loop instead of once per quantization group. + * + * The per-group epilogue term factors exactly: + * + * out[m][n] = ifs[m] * SUM_g wsc[g][n] * ( acc[g] - izp[m]*wsum[g][n] ) + * = ifs[m] * [ SUM_g wsc[g][n]*acc[g] - izp[m]*SUM_g wsc[g][n]*wsum[g][n] ] + * \__ weight-side only __/ + * + * `ifs` is per-row and group-independent so it factors out entirely, and the + * zero-point term separates into a per-row scalar times a per-output-channel + * weight-side sum. That sum depends on no activation data, so it is + * accumulated once into `wcorr_sh` in the prologue rather than being rebuilt + * per group. + * + * Consequences vs dbuf4, per accumulator tile per group: + * - gone: izp*wsum multiply and the subtract (48 v_sub* in the loop body) + * - gone: ifs*wsc multiply (part of 48 dequant-fp) + * - gone: the wsum_sh shared array and its ping-pong + * - gone: izp_bcast / ifs_bcast live across the loop (register pressure) + * - kept: result += float(acc) * wsc + * + * Exact in exact arithmetic, but NOT bit-exact in fp32 -- the summation order + * changes -- so it is gated on the correctness matrix like any other change. + * + * No new binding and no export-format change: the weight-side sum is derived + * in the prologue from t_weight_scales and t_weight_sums, both already bound. + * + * Additionally widens int4 -> int8 byte-parallel (see widen_nibbles below), + * replacing the per-nibble shift/mask/bias-subtract chain. Bit-identical. + * + * Additionally templates the B LDS skew (B_SKEW) so a power-of-two stride can + * be measured against the baseline +1 skew. + * + * Selected via + * ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4zpb_txkgs<32|64>. + * + * Original dbuf4 header follows. + * + * TILE/SUBGROUP-SWEEP variant of the int8 dq8ca_q4gsw coopmat shader's dbuf4 + * ("store-first-for-next", the ORIGINAL loop structure before specs/025 User + * Story 1 picked dbuf2) loop structure (specs/041-dbuf4-tile-sweep). Forked + * from linear_dq8ca_q4gsw_coopmat_tsweep.glsl (which carries dbuf2's loop, + * the production winner) -- everything except the PROLOGUE/MAIN LOOP block + * is identical: bindings, spec-constants, tile-geometry templating, LDS + * layout (ColumnMajor B + skew), int8 WMMA thread maps, group epilog, + * bias/store epilogue. Only the loop structure is swapped to dbuf4, + * recovered from git commit 8d0f23ee78's + * linear_dq8ca_q4gsw_coopmat_dbuf4.glsl (see specs/041/reference/) -- the + * byte-identical pre-swap copy of what is now linear_dq8ca_qw_coopmat.glsl. + * + * The nested `groups x chunks` loop and unconditional group epilog are kept + * exactly as in dbuf2 -- flattening them crashes the Xclipse PAL compiler at + * large spec-resolved trip counts (see dbuf2's own header). Only the + * store/barrier/prefetch ORDER within each chunk iteration is inverted: + * + * dbuf2 (this file's base): store(temp, already prefetched -> cur slice) + * -> barrier -> MMA(cur) -> prefetch(next -> temp) [store owns the + * CURRENT chunk, at the iteration's start] + * dbuf4 (this file): barrier -> prefetch(next -> temp) -> MMA(cur) -> + * store(temp -> next slice) [store owns the NEXT chunk, at the + * iteration's end -- the mirror image] + * + * The group wsum/wsc ping-pong is inverted the same way: dbuf4 stores the + * next group's values (prefetched during the crossing chunk) at the TAIL of + * that chunk, instead of dbuf2's HEAD-of-new-group placement. + * + * Selected at dispatch via + * ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4_txkgs<32|64> + * (QuantizedLinear.cpp), additive to the existing tsweep_t... (dbuf2) + * namespace. + * + * KHR Cooperative Matrix variant of the dynamically-quantized-activation + * linear tiled shader (WEIGHT_NBITS=4): + * 4 -> linear_dq8ca_q4gsw_coopmat INT4 group-symmetric weight + * + * Performs: out[M,N] = dequant(int8_act) * dequant(int_w) (+ bias) + * via coopmat x coopmat -> coopmat on the matrix unit. + * + * Hard preconditions: + * M % WG_TILE_M == 0, N % WG_TILE_N == 0, K % WG_TILE_K == 0, + * INT4: group_size % WG_TILE_K == 0, + * device exposes coopmatx-> at 16x16x16. + */ + +#version 450 core + +#extension GL_KHR_cooperative_matrix : require +#extension GL_KHR_memory_scope_semantics : require +#extension GL_KHR_shader_subgroup_basic : enable +#extension GL_EXT_shader_explicit_arithmetic_types : require +#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_EXT_control_flow_attributes : enable + +#define PRECISION ${PRECISION} + +$if WEIGHT_NBITS == 4: + #define WEIGHT_INT4 + +$if HAS_BIAS: + #define HAS_BIAS + +$if WEIGHT_STORAGE == "buffer": + #define WEIGHT_BUFFER + +$if IO_STORAGE == "texture3d": + #define IO_TEXTURE + +layout(std430) buffer; + +#include "common.glslh" + +// Bindings — match add_linear_dqa_qw_node arg order: +// output(0), fp_input(1), packed_int8_input(2), int_input_sums(3 - unused), +// input_scales(4), input_zps(5), packed_weight(6), weight_sums(7), +// weight_scales(8), bias(9). +${layout_declare_tensor(B, "w", "t_output", "half", IO_STORAGE, is_scalar_array=True)} +// t_input is unread here -- the activations arrive already quantized in +// t_packed_int8_input -- but stays declared so the binding layout matches the +// dispatch site. It tracks IO_STORAGE so the two IO tensors stay consistent. +${layout_declare_tensor(B, "r", "t_input", "half", IO_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_packed_int8_input", "int", "buffer", is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_int8_input_sums", "int", "buffer", is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_int8_input_scales", "half", "texture3d")} +${layout_declare_tensor(B, "r", "t_int8_input_zps", "int8", "texture3d")} +${layout_declare_tensor(B, "r", "t_packed_weight", "int", WEIGHT_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_weight_sums", "int", "buffer", is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_weight_scales", "half", "buffer", is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_bias", "half", "buffer", is_scalar_array=True)} + +${layout_declare_ubo(B, "ivec4", "output_sizes")} +${layout_declare_ubo(B, "ivec4", "input_sizes")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +${layout_declare_spec_const(C, "int", "apply_bias", "0")} +// INT4 only; inert (0) for INT8 so the dispatcher's spec list lines up. +${layout_declare_spec_const(C, "int", "K4_per_group", "0")} +${layout_declare_spec_const(C, "int", "num_groups_arg", "0")} +${layout_declare_spec_const(C, "int", "out_N_arg", "0")} + +// Tile geometry +const uint MMA_M = ${MMA_M}; +const uint MMA_N = ${MMA_N}; +const uint MMA_K = ${MMA_K}; + +const uint WG_TILE_M = ${WG_TILE_M}; +const uint WG_TILE_N = ${WG_TILE_N}; +const uint WG_TILE_K = ${WG_TILE_K}; + +const uint SG_GRID_X = ${SG_GRID_X}; +const uint SG_GRID_Y = ${SG_GRID_Y}; +const uint SUBGROUP_SIZE = ${SUBGROUP_SIZE}; +const uint NUM_SUBGROUPS = SG_GRID_X * SG_GRID_Y; +const uint WG_SIZE = NUM_SUBGROUPS * SUBGROUP_SIZE; + +const uint SG_TILE_M = WG_TILE_M / SG_GRID_Y; +const uint SG_TILE_N = WG_TILE_N / SG_GRID_X; +const uint MMAS_PER_SG_M = SG_TILE_M / MMA_M; +const uint MMAS_PER_SG_N = SG_TILE_N / MMA_N; + +const uint A_SLAB_INT8 = WG_TILE_M * MMA_K; +const uint B_USEFUL_U32 = MMA_K / 4u; +// Intervention A: the bank-conflict skew is what makes B's LDS strides +// non-power-of-two (B_USEFUL_U32+1 = 5, and B_SLAB_U32 = WG_TILE_N*5), forcing +// integer multiplies in every B address. B_SKEW is templated so skew=4 (round +// up to a power of two) and skew=0 (no skew, like the reference -tr's +// ROW_PAD_SH=0) can both be measured against the baseline skew of 1. +const uint B_STRIDE_U32 = B_USEFUL_U32 + ${B_SKEW}u; +const uint B_SLAB_U32 = WG_TILE_N * B_STRIDE_U32; +const uint NUM_K_SLABS = WG_TILE_K / MMA_K; + +const uint A_SLAB_U32 = A_SLAB_INT8 / 4u; +const uint A_STRIDE_U32 = MMA_K / 4u; + +// One ping-pong slice covers all K-slabs of one chunk. +const uint ASH_SLICE_U32 = NUM_K_SLABS * A_SLAB_U32; +const uint BSH_SLICE_U32 = NUM_K_SLABS * B_SLAB_U32; + +// Double-buffered MMA operand staging. +shared uint Ash_int8[2u * ASH_SLICE_U32]; +shared uint Bsh_int8[2u * BSH_SLICE_U32]; + +// Per-WG-tile-row activation params (loaded ONCE at WG start; constant +// across groups). +shared int izp_sh[WG_TILE_M]; // int32 (cast from int8 source) for broadcast +shared float ifs_sh[WG_TILE_M]; // float32 (cast from fp16 source) for broadcast + +// Per-(group, output-channel) weight params, ping-ponged by group parity. +// (For per-channel INT8 only slice 0 is ever used.) +shared float wsc_sh[2u * WG_TILE_N]; +// SUM_g wsc[g][n]*wsum[g][n] per output channel -- weight-side only, so it is +// accumulated once in the prologue. Replaces dbuf4's ping-ponged wsum_sh +// (which was 2*WG_TILE_N ints), so this is a net LDS saving. +shared float wcorr_sh[WG_TILE_N]; + +#ifdef HAS_BIAS +shared float bias_sh[WG_TILE_N]; +#endif + +#ifdef IO_TEXTURE +// Result staging for the imageStore epilogue, mirroring the fp16 kernel: +// SG_GRID_Y bands of MMA_M rows, each WG_TILE_N wide, row-major. A full +// WG_TILE_M x WG_TILE_N buffer would cost SG_GRID_Y/MMAS_PER_SG_M x more LDS +// and wreck occupancy. float16_t-typed because coopMatStore needs it. +const uint CSH_ROWS = SG_GRID_Y * MMA_M; +shared float16_t Csh_out[CSH_ROWS * WG_TILE_N]; +#endif + +// Running fp32 accumulator (across all groups). +coopmat + result[MMAS_PER_SG_M][MMAS_PER_SG_N]; + +// Per-group int32 MMA accumulator. +coopmat + accum_int32[MMAS_PER_SG_M][MMAS_PER_SG_N]; + + +// Byte-parallel int4 -> int8 widening. +// +// The four nibbles this shader needs from one packed uint are ALREADY one per +// byte (bits 3:0 of each byte for parity 0, bits 7:4 for parity 1), so all four +// can be widened at once instead of with a per-nibble +// shift/mask/bias-subtract/mask chain. +// +// For v in [0,15] the biased value is v-8. `v ^ 8` is exactly the 4-bit +// two's-complement pattern of v-8, because -8 == +8 (mod 16): +// v=0 -> 0x8 -> -8 v=7 -> 0xF -> -1 +// v=8 -> 0x0 -> 0 v=15 -> 0x7 -> +7 +// so the only remaining work is sign-extending bit 3 into bits 7:4 per byte. +// `sgn * 0x1E` does that with no cross-byte carry: 0x08 * 0x1E == 0xF0 exactly, +// and sgn is at most 0x08080808 so the product is at most 0xF0F0F0F0. +// +// A naive `nib - 0x08080808` would NOT work -- it borrows across byte lanes +// whenever a nibble is < 8. Shifts must be on uint, not int, so the >> is +// logical rather than arithmetic. +// +// ~5 ops per 4 weights vs ~22 for the per-nibble chain; bit-identical output. +uint widen_nibbles(const uint w, const uint parity) { + const uint nib = (parity == 0u) ? (w & 0x0F0F0F0Fu) : ((w >> 4u) & 0x0F0F0F0Fu); + const uint p = nib ^ 0x08080808u; + const uint sgn = p & 0x08080808u; + return p | (sgn * 0x1Eu); +} + +void main() { + const uvec2 tileID = uvec2(gl_WorkGroupID.xy); + const uvec2 warpInTile = uvec2( + gl_SubgroupID % SG_GRID_X, + gl_SubgroupID / SG_GRID_X); + + const uint K = uint(input_sizes.x); + const uint N = uint(output_sizes.x); + const uint N4 = (N + 3u) / 4u; + const uint nblocks_x_A = (K + 3u) >> 2u; + +#ifdef WEIGHT_INT4 + const uint num_groups = uint(num_groups_arg); + const uint CHUNKS_PER_GROUP = uint(K4_per_group) * 4u / WG_TILE_K; +#else + // Per-channel: a single quant "group" spanning all of K. The nested + // groups x chunks loop below collapses to a flat chunk loop, the wsum/wsc + // ping-pong never crosses a boundary, and the epilog runs exactly once. + const uint num_groups = 1u; + const uint CHUNKS_PER_GROUP = uint(num_groups_arg); +#endif + const uint num_chunks = num_groups * CHUNKS_PER_GROUP; + + const uint tile_m_start = WG_TILE_M * tileID.y; + const uint tile_n_start = WG_TILE_N * tileID.x; + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + result[i][j] = coopmat(0.0); + accum_int32[i][j] = coopmat(0); + } + } + + // --- A staging thread map: one (m4, k4) ivec4 block per active thread --- + const uint K_BLOCKS_PER_CHUNK = WG_TILE_K >> 2u; + const uint A_ACTIVE_THREADS = (WG_TILE_M >> 2u) * K_BLOCKS_PER_CHUNK; + const uint a_m_block = gl_LocalInvocationID.x / K_BLOCKS_PER_CHUNK; + const uint a_k_block = gl_LocalInvocationID.x % K_BLOCKS_PER_CHUNK; + const bool a_active = gl_LocalInvocationID.x < A_ACTIVE_THREADS; + +#ifdef WEIGHT_INT4 + // --- B staging thread map: (block, col) slots; each slot extracts one + // ColumnMajor LDS uint (4 K-contiguous sign-extended int8) --- + const uint B_TOTAL_SLOTS = K_BLOCKS_PER_CHUNK * WG_TILE_N; + const uint B_SLOTS_PER_THREAD = B_TOTAL_SLOTS / WG_SIZE; + const uint N8_PER_TILE = WG_TILE_N >> 3u; +#else + // --- B staging thread map: one (k4, n4) ivec4 block per active thread --- + const uint B_FETCH_SLOTS = K_BLOCKS_PER_CHUNK * (WG_TILE_N >> 2u); + const uint N4_PER_TILE = WG_TILE_N >> 2u; + const uint b_k4_in_chunk = gl_LocalInvocationID.x / N4_PER_TILE; + const uint b_n_uint_col = gl_LocalInvocationID.x % N4_PER_TILE; + const bool b_active = gl_LocalInvocationID.x < B_FETCH_SLOTS; +#endif + + // Prefetch temp registers. + ivec4 temp_A; +#ifdef WEIGHT_INT4 + ivec4 temp_B[B_SLOTS_PER_THREAD]; + float temp_wsc; +#else + ivec4 temp_B; +#endif + + // ========================================================= + // PROLOGUE + // ========================================================= + if (gl_LocalInvocationID.x < (WG_TILE_M >> 2u)) { + const uint m4 = (tile_m_start >> 2u) + gl_LocalInvocationID.x; + const vec4 sc = vec4(texelFetch(t_int8_input_scales, ivec3(m4, 0, 0), 0)); + const ivec4 zp = texelFetch(t_int8_input_zps, ivec3(m4, 0, 0), 0); + const uint base = gl_LocalInvocationID.x * 4u; + ifs_sh[base + 0u] = sc.x; ifs_sh[base + 1u] = sc.y; + ifs_sh[base + 2u] = sc.z; ifs_sh[base + 3u] = sc.w; + izp_sh[base + 0u] = zp.x; izp_sh[base + 1u] = zp.y; + izp_sh[base + 2u] = zp.z; izp_sh[base + 3u] = zp.w; + } + // Group 0 weight scales -> slice 0, and the hoisted weight-side correction + // SUM_g wsc[g][n]*wsum[g][n] accumulated across ALL groups. The loop is + // prologue-only (the prologue is ~1.4% of the dynamic instruction stream), + // and it replaces per-group wsum work inside the loop body. + if (gl_LocalInvocationID.x < WG_TILE_N) { + const uint n_idx = tile_n_start + gl_LocalInvocationID.x; + f16vec4 sv0 = t_weight_scales[n_idx >> 2u]; + wsc_sh[gl_LocalInvocationID.x] = float(sv0[n_idx & 3u]); + + float corr = 0.0; + for (uint g = 0; g < num_groups; ++g) { + f16vec4 sv = t_weight_scales[g * N4 + (n_idx >> 2u)]; + corr += float(sv[n_idx & 3u]) * float(t_weight_sums[g * N + n_idx]); + } + wcorr_sh[gl_LocalInvocationID.x] = corr; + } + memoryBarrierShared(); + barrier(); + + // NOTE: dbuf4 builds izp_bcast/ifs_bcast here and keeps them live across the + // whole group loop. This variant needs them only AFTER the loop, so they are + // loaded there instead -- that is the register-pressure saving. + + // dbuf4: prefetch chunk 0 into temp registers, THEN store to slice 0 (no + // barrier here -- the main loop's first iteration barriers before + // reading slice 0). + if (a_active) { + const uint m4_global = (tile_m_start >> 2u) + a_m_block; + temp_A = t_packed_int8_input[m4_global * nblocks_x_A + a_k_block]; + } +#ifdef WEIGHT_INT4 + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { + const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; + const uint block_in_chunk = slot >> 3u; + const uint k4_blk = block_in_chunk / N8_PER_TILE; + const uint n8_blk = (tile_n_start >> 3u) + (block_in_chunk % N8_PER_TILE); +#ifdef WEIGHT_BUFFER + temp_B[si] = t_packed_weight[(n8_blk * nblocks_x_A) + k4_blk]; +#else + temp_B[si] = texelFetch(t_packed_weight, ivec2(k4_blk, n8_blk), 0); +#endif + } +#else + if (b_active) { + const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; +#ifdef WEIGHT_BUFFER + temp_B = t_packed_weight[(b_k4_in_chunk * N4) + block_x_w]; +#else + temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, b_k4_in_chunk), 0); +#endif + } +#endif + { + // store chunk 0 -> slice 0 + if (a_active) { + const uint slab_idx = a_k_block / (MMA_K >> 2u); + const uint k_uint_in_slab = a_k_block % (MMA_K >> 2u); + const uint base_row = a_m_block * 4u; + [[unroll]] for (uint m4i = 0; m4i < 4u; ++m4i) { + Ash_int8[slab_idx * A_SLAB_U32 + (base_row + m4i) * A_STRIDE_U32 + k_uint_in_slab] = + uint(temp_A[m4i]); + } + } +#ifdef WEIGHT_INT4 + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { + const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; + const uint block_in_chunk = slot >> 3u; + const uint col_in_block = slot & 7u; + const uint k4_in_chunk = block_in_chunk / N8_PER_TILE; + const uint n8_in_tile = block_in_chunk % N8_PER_TILE; + const uint r = col_in_block & 3u; + const uint parity = col_in_block >> 2u; + const uint n_col = n8_in_tile * 8u + r + parity * 4u; + const uint slab_idx = k4_in_chunk / (MMA_K >> 2u); + const uint k4_in_slab = k4_in_chunk % (MMA_K >> 2u); + Bsh_int8[slab_idx * B_SLAB_U32 + n_col * B_STRIDE_U32 + k4_in_slab] = + widen_nibbles(uint(temp_B[si][r]), parity); + } +#else + if (b_active) { + const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); + const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); + const uint n_col_base = b_n_uint_col * 4u; + [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { + Bsh_int8[slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab] = + uint(temp_B[n_in_blk]); + } + } +#endif + } + + // ========================================================= + // MAIN LOOP (dbuf4) — nested groups x chunks (kept nested; flattening it + // with a conditional coopmat epilog crashes the Xclipse PAL compiler at + // large spec-resolved trip counts). One barrier per chunk. Chunk + // iteration (global index `chunk`): + // 1. barrier — A/B slice (chunk%2) fully written; on the first chunk + // of group g, wsum/wsc slice (g%2) is too. + // 2. prefetch — chunk+1 (A blocks, B blocks) into temp; when chunk+1 + // starts a new group, also its wsum/wsc element. Skipped + // entirely on the final chunk. + // 3. int8 MMA — on slice (chunk%2) into accum_int32. + // 4. store — temp -> A/B slice ((chunk+1)%2), unpacking the weight; + // on a group boundary, wsum/wsc -> slice ((g+1)%2). + // The group epilog runs unconditionally at the tail of each group. + // ========================================================= + uint chunk = 0; + for (uint group_i = 0; group_i < num_groups; ++group_i) { + for (uint inner = 0; inner < CHUNKS_PER_GROUP; ++inner, ++chunk) { + const bool has_next = chunk + 1u < num_chunks; + const bool group_crossing = has_next && (inner + 1u == CHUNKS_PER_GROUP); + const uint cur_a = (chunk % 2u) * ASH_SLICE_U32; + const uint cur_b = (chunk % 2u) * BSH_SLICE_U32; + const uint nxt_a = ((chunk + 1u) % 2u) * ASH_SLICE_U32; + const uint nxt_b = ((chunk + 1u) % 2u) * BSH_SLICE_U32; + + barrier(); + + // --- 2. prefetch chunk+1 -> temp --- + if (has_next) { + const uint chunkK_nxt = (chunk + 1u) * WG_TILE_K; + if (a_active) { + const uint m4_global = (tile_m_start >> 2u) + a_m_block; + const uint k4_global = (chunkK_nxt >> 2u) + a_k_block; + temp_A = t_packed_int8_input[m4_global * nblocks_x_A + k4_global]; + } +#ifdef WEIGHT_INT4 + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { + const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; + const uint block_in_chunk = slot >> 3u; + const uint k4_blk = (chunkK_nxt >> 2u) + block_in_chunk / N8_PER_TILE; + const uint n8_blk = (tile_n_start >> 3u) + (block_in_chunk % N8_PER_TILE); +#ifdef WEIGHT_BUFFER + temp_B[si] = t_packed_weight[(n8_blk * nblocks_x_A) + k4_blk]; +#else + temp_B[si] = texelFetch(t_packed_weight, ivec2(k4_blk, n8_blk), 0); +#endif + } + if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { + const uint n_idx = tile_n_start + gl_LocalInvocationID.x; + f16vec4 sv = t_weight_scales[(group_i + 1u) * N4 + (n_idx >> 2u)]; + temp_wsc = float(sv[n_idx & 3u]); + } +#else + if (b_active) { + const uint block_y_w = (chunkK_nxt >> 2u) + b_k4_in_chunk; + const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; +#ifdef WEIGHT_BUFFER + temp_B = t_packed_weight[(block_y_w * N4) + block_x_w]; +#else + temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, block_y_w), 0); +#endif + } +#endif + } + + // --- 3. int8 MMA on the cur slice --- + [[unroll]] for (uint k = 0; k < NUM_K_SLABS; ++k) { + const uint slab_a_base_u32 = cur_a + k * A_SLAB_U32; + const uint slab_b_base_u32 = cur_b + k * B_SLAB_U32; + + coopmat matA[MMAS_PER_SG_M]; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + const uint row_a = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + coopMatLoad( + matA[i], Ash_int8, + slab_a_base_u32 + row_a * A_STRIDE_U32, + A_STRIDE_U32, + gl_CooperativeMatrixLayoutRowMajor); + } + + coopmat matB; + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint col_b = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopMatLoad( + matB, Bsh_int8, + slab_b_base_u32 + col_b * B_STRIDE_U32, + B_STRIDE_U32, + gl_CooperativeMatrixLayoutColumnMajor); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + accum_int32[i][j] = coopMatMulAdd(matA[i], matB, accum_int32[i][j]); + } + } + } + + // --- 4. store temp (chunk+1) -> nxt slice --- + if (has_next) { + if (a_active) { + const uint slab_idx = a_k_block / (MMA_K >> 2u); + const uint k_uint_in_slab = a_k_block % (MMA_K >> 2u); + const uint base_row = a_m_block * 4u; + [[unroll]] for (uint m4i = 0; m4i < 4u; ++m4i) { + Ash_int8[nxt_a + slab_idx * A_SLAB_U32 + (base_row + m4i) * A_STRIDE_U32 + k_uint_in_slab] = + uint(temp_A[m4i]); + } + } +#ifdef WEIGHT_INT4 + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { + const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; + const uint block_in_chunk = slot >> 3u; + const uint col_in_block = slot & 7u; + const uint k4_in_chunk = block_in_chunk / N8_PER_TILE; + const uint n8_in_tile = block_in_chunk % N8_PER_TILE; + const uint r = col_in_block & 3u; + const uint parity = col_in_block >> 2u; + const uint n_col = n8_in_tile * 8u + r + parity * 4u; + const uint slab_idx = k4_in_chunk / (MMA_K >> 2u); + const uint k4_in_slab = k4_in_chunk % (MMA_K >> 2u); + Bsh_int8[nxt_b + slab_idx * B_SLAB_U32 + n_col * B_STRIDE_U32 + k4_in_slab] = + widen_nibbles(uint(temp_B[si][r]), parity); + } + if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { + const uint wbase_nxt = ((group_i + 1u) % 2u) * WG_TILE_N; + wsc_sh[wbase_nxt + gl_LocalInvocationID.x] = temp_wsc; + } +#else + if (b_active) { + const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); + const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); + const uint n_col_base = b_n_uint_col * 4u; + [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { + Bsh_int8[nxt_b + slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab] = + uint(temp_B[n_in_blk]); + } + } +#endif + } + } // chunks + + // --- Group epilog: scale-only accumulate, reset accum --- + // Just result += float(acc) * wsc. The zero-point subtract and the ifs + // multiply are hoisted out of the group loop (applied once below). + { + const uint wbase = (group_i % 2u) * WG_TILE_N; + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + + coopmat wsc_bcast; + coopMatLoad( + wsc_bcast, wsc_sh, + wbase + local_n_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutRowMajor); + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + result[i][j] += + coopmat( + accum_int32[i][j]) * wsc_bcast; + accum_int32[i][j] = coopmat(0); + } + } + } + } // groups + + // --- Hoisted correction, applied ONCE: --------------------------------- + // result = ifs * ( result - izp * SUM_g wsc_g*wsum_g ) + // izp/ifs are loaded here rather than before the group loop so they are not + // live across it. + { + coopmat + izpf_bcast[MMAS_PER_SG_M]; + coopmat + ifs_bcast[MMAS_PER_SG_M]; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + const uint local_m_base = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + coopmat izp_i; + coopMatLoad( + izp_i, izp_sh, local_m_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutColumnMajor); + izpf_bcast[i] = + coopmat(izp_i); + coopMatLoad( + ifs_bcast[i], ifs_sh, local_m_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutColumnMajor); + } + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopmat wcorr_bcast; + coopMatLoad( + wcorr_bcast, wcorr_sh, local_n_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutRowMajor); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + result[i][j] = ifs_bcast[i] * (result[i][j] - izpf_bcast[i] * wcorr_bcast); + } + } + } + + // --- Bias (optional) --- +#ifdef HAS_BIAS + if (apply_bias > 0) { + for (uint t = gl_LocalInvocationID.x; t < WG_TILE_N; t += WG_SIZE) { + bias_sh[t] = float(t_bias[tile_n_start + t]); + } + memoryBarrierShared(); + barrier(); + } +#endif + + // --- Store result tile --- + // N for the store address math MUST come from the spec constant, not the + // sizes UBO (see out_N_arg above). +#ifdef IO_TEXTURE + // Epilogue iteration i drains accumulator row-block i from EVERY subgroup + // into Csh_out at once, so the SG_GRID_Y bands it holds are disjoint global + // row ranges; the whole workgroup then imageStores them. lr / MMA_M is the + // writing subgroup's warpInTile.y, so the global row reproduces the buffer + // path's gi exactly. + // + // PORTABILITY NOTE: the barrier() in the loop body keeps this loop rolled + // despite [[unroll]], so result[i][j] IS dynamically indexed. Coopmat arrays + // are opaque per-lane storage and dynamic indexing is exactly the construct + // the Xclipse/AMD-PAL compiler has broken before -- check this first if the + // texture variants miscompile on M51. + const uint CSH_TEXELS_PER_ROW = WG_TILE_N / 4u; + const uint CSH_TEXELS = CSH_ROWS * CSH_TEXELS_PER_ROW; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + // Guards Csh_out against the previous iteration's readers. Inert on i == 0 + // but must stay unconditional to remain workgroup-uniform. + barrier(); + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { +#ifdef HAS_BIAS + if (apply_bias > 0) { + const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopmat bias_tile; + coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); + result[i][j] += bias_tile; + } +#endif + coopmat out_tile = + coopmat(result[i][j]); + coopMatStore( + out_tile, Csh_out, + warpInTile.y * MMA_M * WG_TILE_N + + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j), + WG_TILE_N, + gl_CooperativeMatrixLayoutRowMajor); + } + memoryBarrierShared(); + barrier(); + + for (uint t = gl_LocalInvocationID.x; t < CSH_TEXELS; t += WG_SIZE) { + const uint lr = t / CSH_TEXELS_PER_ROW; + const uint lc4 = t % CSH_TEXELS_PER_ROW; + const uint m = + tile_m_start + (lr / MMA_M) * SG_TILE_M + i * MMA_M + (lr % MMA_M); + const uint base = lr * WG_TILE_N + lc4 * 4u; + imageStore( + t_output, + ivec3(tile_n_start / 4u + lc4, m, 0), + vec4( + float(Csh_out[base]), + float(Csh_out[base + 1u]), + float(Csh_out[base + 2u]), + float(Csh_out[base + 3u]))); + } + } +#else + const uint N_out = uint(out_N_arg); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint gi = tile_m_start + MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + const uint gj = tile_n_start + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + +#ifdef HAS_BIAS + if (apply_bias > 0) { + const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopmat bias_tile; + coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); + result[i][j] += bias_tile; + } +#endif + + coopmat out_tile = + coopmat(result[i][j]); + coopMatStore( + out_tile, t_output, + gi * N_out + gj, N_out, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#endif // IO_TEXTURE +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpb.yaml b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpb.yaml new file mode 100644 index 00000000000..b2214acce32 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpb.yaml @@ -0,0 +1,42 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# Intervention A: B LDS skew swept over {1 (baseline), 4 (power-of-two stride), +# 0 (no skew)} on top of zp-hoist + byte-parallel nibble widening, at the +# winning tile. Measured target is only 12 v_mul_lo_u32 (1.63% of the loop +# body) and it trades against an LDS category at 8.8%, so a regression is a +# plausible outcome. + +linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpb: + parameter_names_with_default_values: + PRECISION: highp + HAS_BIAS: false + IO_STORAGE: buffer + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + MMA_M: 16 + MMA_N: 16 + MMA_K: 16 + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + B_SKEW: 1 + shader_variants: + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpb_t128x64k32g42s32sk4_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + B_SKEW: 4 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpb_t128x64k32g42s32sk0_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + B_SKEW: 0 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpb_t128x64k32g42s32sk1_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + B_SKEW: 1 diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpn.glsl b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpn.glsl new file mode 100644 index 00000000000..42a813df77d --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpn.glsl @@ -0,0 +1,706 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +/* + * "zp-hoisted" variant: identical to dbuf4 except that the activation + * zero-point correction and the per-row activation scale are applied ONCE + * after the group loop instead of once per quantization group. + * + * The per-group epilogue term factors exactly: + * + * out[m][n] = ifs[m] * SUM_g wsc[g][n] * ( acc[g] - izp[m]*wsum[g][n] ) + * = ifs[m] * [ SUM_g wsc[g][n]*acc[g] - izp[m]*SUM_g wsc[g][n]*wsum[g][n] ] + * \__ weight-side only __/ + * + * `ifs` is per-row and group-independent so it factors out entirely, and the + * zero-point term separates into a per-row scalar times a per-output-channel + * weight-side sum. That sum depends on no activation data, so it is + * accumulated once into `wcorr_sh` in the prologue rather than being rebuilt + * per group. + * + * Consequences vs dbuf4, per accumulator tile per group: + * - gone: izp*wsum multiply and the subtract (48 v_sub* in the loop body) + * - gone: ifs*wsc multiply (part of 48 dequant-fp) + * - gone: the wsum_sh shared array and its ping-pong + * - gone: izp_bcast / ifs_bcast live across the loop (register pressure) + * - kept: result += float(acc) * wsc + * + * Exact in exact arithmetic, but NOT bit-exact in fp32 -- the summation order + * changes -- so it is gated on the correctness matrix like any other change. + * + * No new binding and no export-format change: the weight-side sum is derived + * in the prologue from t_weight_scales and t_weight_sums, both already bound. + * + * Additionally widens int4 -> int8 byte-parallel (see widen_nibbles below), + * replacing the per-nibble shift/mask/bias-subtract chain. Bit-identical. + * + * Selected via + * ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4zpn_txkgs<32|64>. + * + * Original dbuf4 header follows. + * + * TILE/SUBGROUP-SWEEP variant of the int8 dq8ca_q4gsw coopmat shader's dbuf4 + * ("store-first-for-next", the ORIGINAL loop structure before specs/025 User + * Story 1 picked dbuf2) loop structure (specs/041-dbuf4-tile-sweep). Forked + * from linear_dq8ca_q4gsw_coopmat_tsweep.glsl (which carries dbuf2's loop, + * the production winner) -- everything except the PROLOGUE/MAIN LOOP block + * is identical: bindings, spec-constants, tile-geometry templating, LDS + * layout (ColumnMajor B + skew), int8 WMMA thread maps, group epilog, + * bias/store epilogue. Only the loop structure is swapped to dbuf4, + * recovered from git commit 8d0f23ee78's + * linear_dq8ca_q4gsw_coopmat_dbuf4.glsl (see specs/041/reference/) -- the + * byte-identical pre-swap copy of what is now linear_dq8ca_qw_coopmat.glsl. + * + * The nested `groups x chunks` loop and unconditional group epilog are kept + * exactly as in dbuf2 -- flattening them crashes the Xclipse PAL compiler at + * large spec-resolved trip counts (see dbuf2's own header). Only the + * store/barrier/prefetch ORDER within each chunk iteration is inverted: + * + * dbuf2 (this file's base): store(temp, already prefetched -> cur slice) + * -> barrier -> MMA(cur) -> prefetch(next -> temp) [store owns the + * CURRENT chunk, at the iteration's start] + * dbuf4 (this file): barrier -> prefetch(next -> temp) -> MMA(cur) -> + * store(temp -> next slice) [store owns the NEXT chunk, at the + * iteration's end -- the mirror image] + * + * The group wsum/wsc ping-pong is inverted the same way: dbuf4 stores the + * next group's values (prefetched during the crossing chunk) at the TAIL of + * that chunk, instead of dbuf2's HEAD-of-new-group placement. + * + * Selected at dispatch via + * ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4_txkgs<32|64> + * (QuantizedLinear.cpp), additive to the existing tsweep_t... (dbuf2) + * namespace. + * + * KHR Cooperative Matrix variant of the dynamically-quantized-activation + * linear tiled shader (WEIGHT_NBITS=4): + * 4 -> linear_dq8ca_q4gsw_coopmat INT4 group-symmetric weight + * + * Performs: out[M,N] = dequant(int8_act) * dequant(int_w) (+ bias) + * via coopmat x coopmat -> coopmat on the matrix unit. + * + * Hard preconditions: + * M % WG_TILE_M == 0, N % WG_TILE_N == 0, K % WG_TILE_K == 0, + * INT4: group_size % WG_TILE_K == 0, + * device exposes coopmatx-> at 16x16x16. + */ + +#version 450 core + +#extension GL_KHR_cooperative_matrix : require +#extension GL_KHR_memory_scope_semantics : require +#extension GL_KHR_shader_subgroup_basic : enable +#extension GL_EXT_shader_explicit_arithmetic_types : require +#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_EXT_control_flow_attributes : enable + +#define PRECISION ${PRECISION} + +$if WEIGHT_NBITS == 4: + #define WEIGHT_INT4 + +$if HAS_BIAS: + #define HAS_BIAS + +$if WEIGHT_STORAGE == "buffer": + #define WEIGHT_BUFFER + +$if IO_STORAGE == "texture3d": + #define IO_TEXTURE + +layout(std430) buffer; + +#include "common.glslh" + +// Bindings — match add_linear_dqa_qw_node arg order: +// output(0), fp_input(1), packed_int8_input(2), int_input_sums(3 - unused), +// input_scales(4), input_zps(5), packed_weight(6), weight_sums(7), +// weight_scales(8), bias(9). +${layout_declare_tensor(B, "w", "t_output", "half", IO_STORAGE, is_scalar_array=True)} +// t_input is unread here -- the activations arrive already quantized in +// t_packed_int8_input -- but stays declared so the binding layout matches the +// dispatch site. It tracks IO_STORAGE so the two IO tensors stay consistent. +${layout_declare_tensor(B, "r", "t_input", "half", IO_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_packed_int8_input", "int", "buffer", is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_int8_input_sums", "int", "buffer", is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_int8_input_scales", "half", "texture3d")} +${layout_declare_tensor(B, "r", "t_int8_input_zps", "int8", "texture3d")} +${layout_declare_tensor(B, "r", "t_packed_weight", "int", WEIGHT_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_weight_sums", "int", "buffer", is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_weight_scales", "half", "buffer", is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_bias", "half", "buffer", is_scalar_array=True)} + +${layout_declare_ubo(B, "ivec4", "output_sizes")} +${layout_declare_ubo(B, "ivec4", "input_sizes")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +${layout_declare_spec_const(C, "int", "apply_bias", "0")} +// INT4 only; inert (0) for INT8 so the dispatcher's spec list lines up. +${layout_declare_spec_const(C, "int", "K4_per_group", "0")} +${layout_declare_spec_const(C, "int", "num_groups_arg", "0")} +${layout_declare_spec_const(C, "int", "out_N_arg", "0")} + +// Tile geometry +const uint MMA_M = ${MMA_M}; +const uint MMA_N = ${MMA_N}; +const uint MMA_K = ${MMA_K}; + +const uint WG_TILE_M = ${WG_TILE_M}; +const uint WG_TILE_N = ${WG_TILE_N}; +const uint WG_TILE_K = ${WG_TILE_K}; + +const uint SG_GRID_X = ${SG_GRID_X}; +const uint SG_GRID_Y = ${SG_GRID_Y}; +const uint SUBGROUP_SIZE = ${SUBGROUP_SIZE}; +const uint NUM_SUBGROUPS = SG_GRID_X * SG_GRID_Y; +const uint WG_SIZE = NUM_SUBGROUPS * SUBGROUP_SIZE; + +const uint SG_TILE_M = WG_TILE_M / SG_GRID_Y; +const uint SG_TILE_N = WG_TILE_N / SG_GRID_X; +const uint MMAS_PER_SG_M = SG_TILE_M / MMA_M; +const uint MMAS_PER_SG_N = SG_TILE_N / MMA_N; + +const uint A_SLAB_INT8 = WG_TILE_M * MMA_K; +const uint B_USEFUL_U32 = MMA_K / 4u; +const uint B_STRIDE_U32 = B_USEFUL_U32 + 1u; // +1 skew +const uint B_SLAB_U32 = WG_TILE_N * B_STRIDE_U32; +const uint NUM_K_SLABS = WG_TILE_K / MMA_K; + +const uint A_SLAB_U32 = A_SLAB_INT8 / 4u; +const uint A_STRIDE_U32 = MMA_K / 4u; + +// One ping-pong slice covers all K-slabs of one chunk. +const uint ASH_SLICE_U32 = NUM_K_SLABS * A_SLAB_U32; +const uint BSH_SLICE_U32 = NUM_K_SLABS * B_SLAB_U32; + +// Double-buffered MMA operand staging. +shared uint Ash_int8[2u * ASH_SLICE_U32]; +shared uint Bsh_int8[2u * BSH_SLICE_U32]; + +// Per-WG-tile-row activation params (loaded ONCE at WG start; constant +// across groups). +shared int izp_sh[WG_TILE_M]; // int32 (cast from int8 source) for broadcast +shared float ifs_sh[WG_TILE_M]; // float32 (cast from fp16 source) for broadcast + +// Per-(group, output-channel) weight params, ping-ponged by group parity. +// (For per-channel INT8 only slice 0 is ever used.) +shared float wsc_sh[2u * WG_TILE_N]; +// SUM_g wsc[g][n]*wsum[g][n] per output channel -- weight-side only, so it is +// accumulated once in the prologue. Replaces dbuf4's ping-ponged wsum_sh +// (which was 2*WG_TILE_N ints), so this is a net LDS saving. +shared float wcorr_sh[WG_TILE_N]; + +#ifdef HAS_BIAS +shared float bias_sh[WG_TILE_N]; +#endif + +#ifdef IO_TEXTURE +// Result staging for the imageStore epilogue, mirroring the fp16 kernel: +// SG_GRID_Y bands of MMA_M rows, each WG_TILE_N wide, row-major. A full +// WG_TILE_M x WG_TILE_N buffer would cost SG_GRID_Y/MMAS_PER_SG_M x more LDS +// and wreck occupancy. float16_t-typed because coopMatStore needs it. +const uint CSH_ROWS = SG_GRID_Y * MMA_M; +shared float16_t Csh_out[CSH_ROWS * WG_TILE_N]; +#endif + +// Running fp32 accumulator (across all groups). +coopmat + result[MMAS_PER_SG_M][MMAS_PER_SG_N]; + +// Per-group int32 MMA accumulator. +coopmat + accum_int32[MMAS_PER_SG_M][MMAS_PER_SG_N]; + + +// Byte-parallel int4 -> int8 widening. +// +// The four nibbles this shader needs from one packed uint are ALREADY one per +// byte (bits 3:0 of each byte for parity 0, bits 7:4 for parity 1), so all four +// can be widened at once instead of with a per-nibble +// shift/mask/bias-subtract/mask chain. +// +// For v in [0,15] the biased value is v-8. `v ^ 8` is exactly the 4-bit +// two's-complement pattern of v-8, because -8 == +8 (mod 16): +// v=0 -> 0x8 -> -8 v=7 -> 0xF -> -1 +// v=8 -> 0x0 -> 0 v=15 -> 0x7 -> +7 +// so the only remaining work is sign-extending bit 3 into bits 7:4 per byte. +// `sgn * 0x1E` does that with no cross-byte carry: 0x08 * 0x1E == 0xF0 exactly, +// and sgn is at most 0x08080808 so the product is at most 0xF0F0F0F0. +// +// A naive `nib - 0x08080808` would NOT work -- it borrows across byte lanes +// whenever a nibble is < 8. Shifts must be on uint, not int, so the >> is +// logical rather than arithmetic. +// +// ~5 ops per 4 weights vs ~22 for the per-nibble chain; bit-identical output. +uint widen_nibbles(const uint w, const uint parity) { + const uint nib = (parity == 0u) ? (w & 0x0F0F0F0Fu) : ((w >> 4u) & 0x0F0F0F0Fu); + const uint p = nib ^ 0x08080808u; + const uint sgn = p & 0x08080808u; + return p | (sgn * 0x1Eu); +} + +void main() { + const uvec2 tileID = uvec2(gl_WorkGroupID.xy); + const uvec2 warpInTile = uvec2( + gl_SubgroupID % SG_GRID_X, + gl_SubgroupID / SG_GRID_X); + + const uint K = uint(input_sizes.x); + const uint N = uint(output_sizes.x); + const uint N4 = (N + 3u) / 4u; + const uint nblocks_x_A = (K + 3u) >> 2u; + +#ifdef WEIGHT_INT4 + const uint num_groups = uint(num_groups_arg); + const uint CHUNKS_PER_GROUP = uint(K4_per_group) * 4u / WG_TILE_K; +#else + // Per-channel: a single quant "group" spanning all of K. The nested + // groups x chunks loop below collapses to a flat chunk loop, the wsum/wsc + // ping-pong never crosses a boundary, and the epilog runs exactly once. + const uint num_groups = 1u; + const uint CHUNKS_PER_GROUP = uint(num_groups_arg); +#endif + const uint num_chunks = num_groups * CHUNKS_PER_GROUP; + + const uint tile_m_start = WG_TILE_M * tileID.y; + const uint tile_n_start = WG_TILE_N * tileID.x; + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + result[i][j] = coopmat(0.0); + accum_int32[i][j] = coopmat(0); + } + } + + // --- A staging thread map: one (m4, k4) ivec4 block per active thread --- + const uint K_BLOCKS_PER_CHUNK = WG_TILE_K >> 2u; + const uint A_ACTIVE_THREADS = (WG_TILE_M >> 2u) * K_BLOCKS_PER_CHUNK; + const uint a_m_block = gl_LocalInvocationID.x / K_BLOCKS_PER_CHUNK; + const uint a_k_block = gl_LocalInvocationID.x % K_BLOCKS_PER_CHUNK; + const bool a_active = gl_LocalInvocationID.x < A_ACTIVE_THREADS; + +#ifdef WEIGHT_INT4 + // --- B staging thread map: (block, col) slots; each slot extracts one + // ColumnMajor LDS uint (4 K-contiguous sign-extended int8) --- + const uint B_TOTAL_SLOTS = K_BLOCKS_PER_CHUNK * WG_TILE_N; + const uint B_SLOTS_PER_THREAD = B_TOTAL_SLOTS / WG_SIZE; + const uint N8_PER_TILE = WG_TILE_N >> 3u; +#else + // --- B staging thread map: one (k4, n4) ivec4 block per active thread --- + const uint B_FETCH_SLOTS = K_BLOCKS_PER_CHUNK * (WG_TILE_N >> 2u); + const uint N4_PER_TILE = WG_TILE_N >> 2u; + const uint b_k4_in_chunk = gl_LocalInvocationID.x / N4_PER_TILE; + const uint b_n_uint_col = gl_LocalInvocationID.x % N4_PER_TILE; + const bool b_active = gl_LocalInvocationID.x < B_FETCH_SLOTS; +#endif + + // Prefetch temp registers. + ivec4 temp_A; +#ifdef WEIGHT_INT4 + ivec4 temp_B[B_SLOTS_PER_THREAD]; + float temp_wsc; +#else + ivec4 temp_B; +#endif + + // ========================================================= + // PROLOGUE + // ========================================================= + if (gl_LocalInvocationID.x < (WG_TILE_M >> 2u)) { + const uint m4 = (tile_m_start >> 2u) + gl_LocalInvocationID.x; + const vec4 sc = vec4(texelFetch(t_int8_input_scales, ivec3(m4, 0, 0), 0)); + const ivec4 zp = texelFetch(t_int8_input_zps, ivec3(m4, 0, 0), 0); + const uint base = gl_LocalInvocationID.x * 4u; + ifs_sh[base + 0u] = sc.x; ifs_sh[base + 1u] = sc.y; + ifs_sh[base + 2u] = sc.z; ifs_sh[base + 3u] = sc.w; + izp_sh[base + 0u] = zp.x; izp_sh[base + 1u] = zp.y; + izp_sh[base + 2u] = zp.z; izp_sh[base + 3u] = zp.w; + } + // Group 0 weight scales -> slice 0, and the hoisted weight-side correction + // SUM_g wsc[g][n]*wsum[g][n] accumulated across ALL groups. The loop is + // prologue-only (the prologue is ~1.4% of the dynamic instruction stream), + // and it replaces per-group wsum work inside the loop body. + if (gl_LocalInvocationID.x < WG_TILE_N) { + const uint n_idx = tile_n_start + gl_LocalInvocationID.x; + f16vec4 sv0 = t_weight_scales[n_idx >> 2u]; + wsc_sh[gl_LocalInvocationID.x] = float(sv0[n_idx & 3u]); + + float corr = 0.0; + for (uint g = 0; g < num_groups; ++g) { + f16vec4 sv = t_weight_scales[g * N4 + (n_idx >> 2u)]; + corr += float(sv[n_idx & 3u]) * float(t_weight_sums[g * N + n_idx]); + } + wcorr_sh[gl_LocalInvocationID.x] = corr; + } + memoryBarrierShared(); + barrier(); + + // NOTE: dbuf4 builds izp_bcast/ifs_bcast here and keeps them live across the + // whole group loop. This variant needs them only AFTER the loop, so they are + // loaded there instead -- that is the register-pressure saving. + + // dbuf4: prefetch chunk 0 into temp registers, THEN store to slice 0 (no + // barrier here -- the main loop's first iteration barriers before + // reading slice 0). + if (a_active) { + const uint m4_global = (tile_m_start >> 2u) + a_m_block; + temp_A = t_packed_int8_input[m4_global * nblocks_x_A + a_k_block]; + } +#ifdef WEIGHT_INT4 + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { + const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; + const uint block_in_chunk = slot >> 3u; + const uint k4_blk = block_in_chunk / N8_PER_TILE; + const uint n8_blk = (tile_n_start >> 3u) + (block_in_chunk % N8_PER_TILE); +#ifdef WEIGHT_BUFFER + temp_B[si] = t_packed_weight[(n8_blk * nblocks_x_A) + k4_blk]; +#else + temp_B[si] = texelFetch(t_packed_weight, ivec2(k4_blk, n8_blk), 0); +#endif + } +#else + if (b_active) { + const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; +#ifdef WEIGHT_BUFFER + temp_B = t_packed_weight[(b_k4_in_chunk * N4) + block_x_w]; +#else + temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, b_k4_in_chunk), 0); +#endif + } +#endif + { + // store chunk 0 -> slice 0 + if (a_active) { + const uint slab_idx = a_k_block / (MMA_K >> 2u); + const uint k_uint_in_slab = a_k_block % (MMA_K >> 2u); + const uint base_row = a_m_block * 4u; + [[unroll]] for (uint m4i = 0; m4i < 4u; ++m4i) { + Ash_int8[slab_idx * A_SLAB_U32 + (base_row + m4i) * A_STRIDE_U32 + k_uint_in_slab] = + uint(temp_A[m4i]); + } + } +#ifdef WEIGHT_INT4 + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { + const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; + const uint block_in_chunk = slot >> 3u; + const uint col_in_block = slot & 7u; + const uint k4_in_chunk = block_in_chunk / N8_PER_TILE; + const uint n8_in_tile = block_in_chunk % N8_PER_TILE; + const uint r = col_in_block & 3u; + const uint parity = col_in_block >> 2u; + const uint n_col = n8_in_tile * 8u + r + parity * 4u; + const uint slab_idx = k4_in_chunk / (MMA_K >> 2u); + const uint k4_in_slab = k4_in_chunk % (MMA_K >> 2u); + Bsh_int8[slab_idx * B_SLAB_U32 + n_col * B_STRIDE_U32 + k4_in_slab] = + widen_nibbles(uint(temp_B[si][r]), parity); + } +#else + if (b_active) { + const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); + const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); + const uint n_col_base = b_n_uint_col * 4u; + [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { + Bsh_int8[slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab] = + uint(temp_B[n_in_blk]); + } + } +#endif + } + + // ========================================================= + // MAIN LOOP (dbuf4) — nested groups x chunks (kept nested; flattening it + // with a conditional coopmat epilog crashes the Xclipse PAL compiler at + // large spec-resolved trip counts). One barrier per chunk. Chunk + // iteration (global index `chunk`): + // 1. barrier — A/B slice (chunk%2) fully written; on the first chunk + // of group g, wsum/wsc slice (g%2) is too. + // 2. prefetch — chunk+1 (A blocks, B blocks) into temp; when chunk+1 + // starts a new group, also its wsum/wsc element. Skipped + // entirely on the final chunk. + // 3. int8 MMA — on slice (chunk%2) into accum_int32. + // 4. store — temp -> A/B slice ((chunk+1)%2), unpacking the weight; + // on a group boundary, wsum/wsc -> slice ((g+1)%2). + // The group epilog runs unconditionally at the tail of each group. + // ========================================================= + uint chunk = 0; + for (uint group_i = 0; group_i < num_groups; ++group_i) { + for (uint inner = 0; inner < CHUNKS_PER_GROUP; ++inner, ++chunk) { + const bool has_next = chunk + 1u < num_chunks; + const bool group_crossing = has_next && (inner + 1u == CHUNKS_PER_GROUP); + const uint cur_a = (chunk % 2u) * ASH_SLICE_U32; + const uint cur_b = (chunk % 2u) * BSH_SLICE_U32; + const uint nxt_a = ((chunk + 1u) % 2u) * ASH_SLICE_U32; + const uint nxt_b = ((chunk + 1u) % 2u) * BSH_SLICE_U32; + + barrier(); + + // --- 2. prefetch chunk+1 -> temp --- + if (has_next) { + const uint chunkK_nxt = (chunk + 1u) * WG_TILE_K; + if (a_active) { + const uint m4_global = (tile_m_start >> 2u) + a_m_block; + const uint k4_global = (chunkK_nxt >> 2u) + a_k_block; + temp_A = t_packed_int8_input[m4_global * nblocks_x_A + k4_global]; + } +#ifdef WEIGHT_INT4 + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { + const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; + const uint block_in_chunk = slot >> 3u; + const uint k4_blk = (chunkK_nxt >> 2u) + block_in_chunk / N8_PER_TILE; + const uint n8_blk = (tile_n_start >> 3u) + (block_in_chunk % N8_PER_TILE); +#ifdef WEIGHT_BUFFER + temp_B[si] = t_packed_weight[(n8_blk * nblocks_x_A) + k4_blk]; +#else + temp_B[si] = texelFetch(t_packed_weight, ivec2(k4_blk, n8_blk), 0); +#endif + } + if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { + const uint n_idx = tile_n_start + gl_LocalInvocationID.x; + f16vec4 sv = t_weight_scales[(group_i + 1u) * N4 + (n_idx >> 2u)]; + temp_wsc = float(sv[n_idx & 3u]); + } +#else + if (b_active) { + const uint block_y_w = (chunkK_nxt >> 2u) + b_k4_in_chunk; + const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; +#ifdef WEIGHT_BUFFER + temp_B = t_packed_weight[(block_y_w * N4) + block_x_w]; +#else + temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, block_y_w), 0); +#endif + } +#endif + } + + // --- 3. int8 MMA on the cur slice --- + [[unroll]] for (uint k = 0; k < NUM_K_SLABS; ++k) { + const uint slab_a_base_u32 = cur_a + k * A_SLAB_U32; + const uint slab_b_base_u32 = cur_b + k * B_SLAB_U32; + + coopmat matA[MMAS_PER_SG_M]; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + const uint row_a = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + coopMatLoad( + matA[i], Ash_int8, + slab_a_base_u32 + row_a * A_STRIDE_U32, + A_STRIDE_U32, + gl_CooperativeMatrixLayoutRowMajor); + } + + coopmat matB; + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint col_b = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopMatLoad( + matB, Bsh_int8, + slab_b_base_u32 + col_b * B_STRIDE_U32, + B_STRIDE_U32, + gl_CooperativeMatrixLayoutColumnMajor); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + accum_int32[i][j] = coopMatMulAdd(matA[i], matB, accum_int32[i][j]); + } + } + } + + // --- 4. store temp (chunk+1) -> nxt slice --- + if (has_next) { + if (a_active) { + const uint slab_idx = a_k_block / (MMA_K >> 2u); + const uint k_uint_in_slab = a_k_block % (MMA_K >> 2u); + const uint base_row = a_m_block * 4u; + [[unroll]] for (uint m4i = 0; m4i < 4u; ++m4i) { + Ash_int8[nxt_a + slab_idx * A_SLAB_U32 + (base_row + m4i) * A_STRIDE_U32 + k_uint_in_slab] = + uint(temp_A[m4i]); + } + } +#ifdef WEIGHT_INT4 + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { + const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; + const uint block_in_chunk = slot >> 3u; + const uint col_in_block = slot & 7u; + const uint k4_in_chunk = block_in_chunk / N8_PER_TILE; + const uint n8_in_tile = block_in_chunk % N8_PER_TILE; + const uint r = col_in_block & 3u; + const uint parity = col_in_block >> 2u; + const uint n_col = n8_in_tile * 8u + r + parity * 4u; + const uint slab_idx = k4_in_chunk / (MMA_K >> 2u); + const uint k4_in_slab = k4_in_chunk % (MMA_K >> 2u); + Bsh_int8[nxt_b + slab_idx * B_SLAB_U32 + n_col * B_STRIDE_U32 + k4_in_slab] = + widen_nibbles(uint(temp_B[si][r]), parity); + } + if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { + const uint wbase_nxt = ((group_i + 1u) % 2u) * WG_TILE_N; + wsc_sh[wbase_nxt + gl_LocalInvocationID.x] = temp_wsc; + } +#else + if (b_active) { + const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); + const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); + const uint n_col_base = b_n_uint_col * 4u; + [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { + Bsh_int8[nxt_b + slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab] = + uint(temp_B[n_in_blk]); + } + } +#endif + } + } // chunks + + // --- Group epilog: scale-only accumulate, reset accum --- + // Just result += float(acc) * wsc. The zero-point subtract and the ifs + // multiply are hoisted out of the group loop (applied once below). + { + const uint wbase = (group_i % 2u) * WG_TILE_N; + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + + coopmat wsc_bcast; + coopMatLoad( + wsc_bcast, wsc_sh, + wbase + local_n_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutRowMajor); + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + result[i][j] += + coopmat( + accum_int32[i][j]) * wsc_bcast; + accum_int32[i][j] = coopmat(0); + } + } + } + } // groups + + // --- Hoisted correction, applied ONCE: --------------------------------- + // result = ifs * ( result - izp * SUM_g wsc_g*wsum_g ) + // izp/ifs are loaded here rather than before the group loop so they are not + // live across it. + { + coopmat + izpf_bcast[MMAS_PER_SG_M]; + coopmat + ifs_bcast[MMAS_PER_SG_M]; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + const uint local_m_base = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + coopmat izp_i; + coopMatLoad( + izp_i, izp_sh, local_m_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutColumnMajor); + izpf_bcast[i] = + coopmat(izp_i); + coopMatLoad( + ifs_bcast[i], ifs_sh, local_m_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutColumnMajor); + } + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopmat wcorr_bcast; + coopMatLoad( + wcorr_bcast, wcorr_sh, local_n_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutRowMajor); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + result[i][j] = ifs_bcast[i] * (result[i][j] - izpf_bcast[i] * wcorr_bcast); + } + } + } + + // --- Bias (optional) --- +#ifdef HAS_BIAS + if (apply_bias > 0) { + for (uint t = gl_LocalInvocationID.x; t < WG_TILE_N; t += WG_SIZE) { + bias_sh[t] = float(t_bias[tile_n_start + t]); + } + memoryBarrierShared(); + barrier(); + } +#endif + + // --- Store result tile --- + // N for the store address math MUST come from the spec constant, not the + // sizes UBO (see out_N_arg above). +#ifdef IO_TEXTURE + // Epilogue iteration i drains accumulator row-block i from EVERY subgroup + // into Csh_out at once, so the SG_GRID_Y bands it holds are disjoint global + // row ranges; the whole workgroup then imageStores them. lr / MMA_M is the + // writing subgroup's warpInTile.y, so the global row reproduces the buffer + // path's gi exactly. + // + // PORTABILITY NOTE: the barrier() in the loop body keeps this loop rolled + // despite [[unroll]], so result[i][j] IS dynamically indexed. Coopmat arrays + // are opaque per-lane storage and dynamic indexing is exactly the construct + // the Xclipse/AMD-PAL compiler has broken before -- check this first if the + // texture variants miscompile on M51. + const uint CSH_TEXELS_PER_ROW = WG_TILE_N / 4u; + const uint CSH_TEXELS = CSH_ROWS * CSH_TEXELS_PER_ROW; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + // Guards Csh_out against the previous iteration's readers. Inert on i == 0 + // but must stay unconditional to remain workgroup-uniform. + barrier(); + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { +#ifdef HAS_BIAS + if (apply_bias > 0) { + const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopmat bias_tile; + coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); + result[i][j] += bias_tile; + } +#endif + coopmat out_tile = + coopmat(result[i][j]); + coopMatStore( + out_tile, Csh_out, + warpInTile.y * MMA_M * WG_TILE_N + + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j), + WG_TILE_N, + gl_CooperativeMatrixLayoutRowMajor); + } + memoryBarrierShared(); + barrier(); + + for (uint t = gl_LocalInvocationID.x; t < CSH_TEXELS; t += WG_SIZE) { + const uint lr = t / CSH_TEXELS_PER_ROW; + const uint lc4 = t % CSH_TEXELS_PER_ROW; + const uint m = + tile_m_start + (lr / MMA_M) * SG_TILE_M + i * MMA_M + (lr % MMA_M); + const uint base = lr * WG_TILE_N + lc4 * 4u; + imageStore( + t_output, + ivec3(tile_n_start / 4u + lc4, m, 0), + vec4( + float(Csh_out[base]), + float(Csh_out[base + 1u]), + float(Csh_out[base + 2u]), + float(Csh_out[base + 3u]))); + } + } +#else + const uint N_out = uint(out_N_arg); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint gi = tile_m_start + MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + const uint gj = tile_n_start + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + +#ifdef HAS_BIAS + if (apply_bias > 0) { + const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopmat bias_tile; + coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); + result[i][j] += bias_tile; + } +#endif + + coopmat out_tile = + coopmat(result[i][j]); + coopMatStore( + out_tile, t_output, + gi * N_out + gj, N_out, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#endif // IO_TEXTURE +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpn.yaml b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpn.yaml new file mode 100644 index 00000000000..b40757f0fc0 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpn.yaml @@ -0,0 +1,187 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# "zp-hoisted" variants of the int8 dq8ca_q4gsw coopmat kernel: identical to +# tsweep_dbuf4 except the activation zero-point correction and the per-row +# activation scale are applied once after the group loop instead of once per +# quantization group. No new binding, no export-format change -- the weight-side +# correction sum is derived in the prologue from tensors already bound. +# +# Selected via +# ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4zpn_txkgs. +# NOT the shipped default until it passes repeated correctness runs. + +linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpn: + parameter_names_with_default_values: + PRECISION: highp + HAS_BIAS: false + IO_STORAGE: buffer + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + MMA_M: 16 + MMA_N: 16 + MMA_K: 16 + WG_TILE_M: 64 + WG_TILE_N: 32 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + shader_variants: + # production tile -- the A/B anchor against the 1.3 baseline + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpn_t64x32k32g12s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 32 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpn_t64x32k32g12s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 32 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpn_t64x32k32g12s64_buffer_buffer_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: buffer + WG_TILE_M: 64 + WG_TILE_N: 32 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + # NOTE: t128x128k64g42s32 was seeded here and then REMOVED. Its A-staging + # thread map cannot cover its blocks: A_ACTIVE_THREADS = (128/4)*(64/4) = + # 512 but WG_SIZE = 4*2*32 = 256, so half of A is never staged and the + # kernel is numerically wrong (measured 12/12 correctness failures on the + # dbuf4 equivalent). Use g44s32 for that tile instead -- 4*4*32 = 512 = A. + # deep K + big tile; the reference's winning geometry, A map valid at 512=512 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpn_t128x128k64g44s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 4 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpn_t128x128k64g44s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 4 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + # task 5.1's explicit ask + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpn_t128x128k32g42s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpn_t128x128k32g42s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + # same tile, more subgroups -> MMAS 2x2 instead of 4x2 (less accumulator pressure) + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpn_t128x128k32g44s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpn_t128x128k32g44s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 4 + SUBGROUP_SIZE: 32 + # moderate step up from the shipped 64x32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpn_t128x64k32g42s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpn_t128x64k32g42s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + # small step up + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpn_t64x64k32g22s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpn_t64x64k32g22s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 64 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + # big tile, shallow K -- isolates tile area from K depth + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpn_t128x128k16g42s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpn_t128x128k16g42s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 16 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp index ab9201d8af8..c224925cd69 100644 --- a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp @@ -100,6 +100,17 @@ static const char* const kTsweepPrefixes[] = { // '_', so prefix order here does not matter. dq8ca only -- there is no // q4gsw "-tr" shader. "tsweep_dbuf4tr_t", + // DIAGNOSTIC: same row-major layout as "-tr" but with scalar A staging, + // used to bisect a dbuf4tr correctness failure. Delete with the shader. + "tsweep_dbuf4trm_t", + "tsweep_dbuf4trd_t", + // zp-hoisted: zero-point + activation-scale correction applied once + // after the group loop instead of once per quantization group. + "tsweep_dbuf4zp_t", + // zp-hoisted + byte-parallel int4->int8 widening. + "tsweep_dbuf4zpn_t", + // + templated B LDS skew (intervention A). + "tsweep_dbuf4zpb_t", "tsweep_t", }; @@ -438,7 +449,10 @@ static bool can_use_q4gsw_coopmat( // because no coopMatLoad can address 4h4w (its component index selects a row, // making the flat index non-affine in the row). static bool dq8ca_variant_wants_rowmajor_a() { - return dq8ca_coopmat_variant().rfind("tsweep_dbuf4tr_t", 0) == 0; + const std::string& v = dq8ca_coopmat_variant(); + return v.rfind("tsweep_dbuf4tr_t", 0) == 0 || + v.rfind("tsweep_dbuf4trm_t", 0) == 0 || + v.rfind("tsweep_dbuf4trd_t", 0) == 0; } // Mirrors the coopmat branch of pick_linear_dqa_qw_shader() so graph-build time From 8cde63eae4a6c28435703c4e4436e2fd30aa7810 Mon Sep 17 00:00:00 2001 From: Yanwen Xu Date: Wed, 26 Aug 2026 14:43:33 -0700 Subject: [PATCH 12/28] [ET-VK] Speed up 8da4w coopmat linear: bigger tile, drop B skew, coalesce B store Three real, on-device-validated changes to the dq8ca (8da4w) coopmat linear kernel, found via the dq8ca-dequant-unpack-ablation investigation into why this kernel trails a dense-int8 reference by a wide margin: - Promote the shipped tile from t64x32k32g12s64 to t128x128k64g81s64 (already ahead pre-fix; +11.8-13.2% relative on top of the other two changes). - Drop the classic "+1" anti-bank-conflict skew on B's LDS stride (B_STRIDE_U32): measured slower than no padding on this hardware/driver (+2.98pp on its own). - Make the B-operand LDS write coalesced: invert the write-address formula so consecutive threads write consecutive LDS words, instead of 4 words apart. The read side and every dequantized value are unchanged; only which thread writes which element (+0.7-0.8% on top of the other two). Combined: 8B prefill efficiency of int8 theoretical peak goes from 28.93% (original) to 36.97%+ (now shipped) on the primary validation board. Real e2e prefill: 1B 1532.9, 3B 616.9, 8B 299.2 tok/s (ETDump-confirmed coopmat dispatch, ET_VK_TEXTURE_COOPMAT=1). Validated to this kernel family's own stated bar: 10 consecutive --correctness-only passes, zero failures, for every (model x storage) combination -- 1B/3B/8B x buffer/texture3d, 60+ runs total. Decode-regime validation and a second board/driver are still open. Full investigation, ablation chain, and the (still open) question of why the B-operand LDS read itself is so much more expensive than a dense-int8 reference's is in openspec/changes/dq8ca-dequant-unpack-ablation/results/README.md. Authored with Claude. --- ...near_dq8ca_q4gsw_coopmat_tsweep_dbuf4.glsl | 106 +++++++++------ .../graph/ops/impl/QuantizedLinear.cpp | 122 ++++++++++++++---- 2 files changed, 167 insertions(+), 61 deletions(-) diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4.glsl b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4.glsl index b8d4b309233..80408194843 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4.glsl @@ -132,7 +132,13 @@ const uint MMAS_PER_SG_N = SG_TILE_N / MMA_N; const uint A_SLAB_INT8 = WG_TILE_M * MMA_K; const uint B_USEFUL_U32 = MMA_K / 4u; -const uint B_STRIDE_U32 = B_USEFUL_U32 + 1u; // +1 skew +// No skew. The classic anti-bank-conflict "+1" padding was measured SLOWER on M51 at the shipped +// tile (t64x32k32g12s64): stride=4 (this) is +2.98pp efficiency over stride=5 (the old +1 skew), +// stride=6/+2 is a wash, stride=8/+4 is worse -- real, on-device, dq8ca-dequant-unpack-ablation +// (openspec/changes/dq8ca-dequant-unpack-ablation/results/README.md, Addendum 6). B_USEFUL_U32 is +// tile-invariant (MMA_K is fixed at 16 for every variant in this family), so this applies uniformly; +// only re-validated at the shipped tile specifically, not every swept variant. +const uint B_STRIDE_U32 = B_USEFUL_U32; const uint B_SLAB_U32 = WG_TILE_N * B_STRIDE_U32; const uint NUM_K_SLABS = WG_TILE_K / MMA_K; @@ -170,6 +176,44 @@ const uint CSH_ROWS = SG_GRID_Y * MMA_M; shared float16_t Csh_out[CSH_ROWS * WG_TILE_N]; #endif +#ifdef WEIGHT_INT4 +// Coalesced-B-write address inversion: given the contiguous per-thread LDS +// index `a` (in [0, BSH_SLICE_U32)) this thread will write to, recover which +// global-fetch element it needs. Consecutive threads (consecutive `a`) now +// write consecutive LDS words -- unlike the pre-2026-08-26 mapping, where 8 +// consecutive threads wrote addresses B_STRIDE_U32 words apart. Real, +// on-device measurement (dq8ca-dequant-unpack-ablation follow-up, +// openspec/changes/dq8ca-dequant-unpack-ablation/results/README.md, "the +// coalesced B-store rewrite"): a consistent +0.7-0.8% real speedup across +// 1B/3B/8B at the shipped tile, correctness-clean (10/10 buffer + 10/10 +// texture3d, all three models). `chunkK_base` is the K-offset of the chunk +// being staged (0 for the prologue, chunkK_nxt for the main loop's chunk+1 +// prefetch; the store sites don't need it -- r/parity depend only on n_col, +// not chunkK_base -- so they pass 0u). +struct BCoalIndex { + uint n8_blk; + uint k4_blk; + uint r; + uint parity; +}; + +BCoalIndex bcoal_index(const uint a, const uint chunkK_base, const uint tile_n_start) { + const uint slab_idx = a / B_SLAB_U32; + const uint local_a = a % B_SLAB_U32; + const uint n_col = local_a / B_STRIDE_U32; + const uint k4_in_slab = local_a % B_STRIDE_U32; + const uint k4_in_chunk = slab_idx * (MMA_K >> 2u) + k4_in_slab; + const uint n8_in_tile = n_col >> 3u; + const uint rem = n_col & 7u; + BCoalIndex idx; + idx.n8_blk = (tile_n_start >> 3u) + n8_in_tile; + idx.k4_blk = (chunkK_base >> 2u) + k4_in_chunk; + idx.r = rem & 3u; + idx.parity = rem >> 2u; + return idx; +} +#endif + // Running fp32 accumulator (across all groups). coopmat result[MMAS_PER_SG_M][MMAS_PER_SG_N]; @@ -294,14 +338,12 @@ void main() { } #ifdef WEIGHT_INT4 [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { - const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; - const uint block_in_chunk = slot >> 3u; - const uint k4_blk = block_in_chunk / N8_PER_TILE; - const uint n8_blk = (tile_n_start >> 3u) + (block_in_chunk % N8_PER_TILE); + const uint a = gl_LocalInvocationID.x + si * WG_SIZE; + const BCoalIndex bidx0 = bcoal_index(a, 0u, tile_n_start); #ifdef WEIGHT_BUFFER - temp_B[si] = t_packed_weight[(n8_blk * nblocks_x_A) + k4_blk]; + temp_B[si] = t_packed_weight[(bidx0.n8_blk * nblocks_x_A) + bidx0.k4_blk]; #else - temp_B[si] = texelFetch(t_packed_weight, ivec2(k4_blk, n8_blk), 0); + temp_B[si] = texelFetch(t_packed_weight, ivec2(bidx0.k4_blk, bidx0.n8_blk), 0); #endif } #else @@ -327,23 +369,15 @@ void main() { } #ifdef WEIGHT_INT4 [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { - const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; - const uint block_in_chunk = slot >> 3u; - const uint col_in_block = slot & 7u; - const uint k4_in_chunk = block_in_chunk / N8_PER_TILE; - const uint n8_in_tile = block_in_chunk % N8_PER_TILE; - const uint r = col_in_block & 3u; - const uint parity = col_in_block >> 2u; - const int w = temp_B[si][r]; - const int base = int(4u * parity); + const uint a = gl_LocalInvocationID.x + si * WG_SIZE; + const BCoalIndex bidx0 = bcoal_index(a, 0u, tile_n_start); + const int w = temp_B[si][bidx0.r]; + const int base = int(4u * bidx0.parity); const int v0 = (((w >> (base + 0)) & 0xF) - 8) & 0xFF; const int v1 = (((w >> (base + 8)) & 0xF) - 8) & 0xFF; const int v2 = (((w >> (base + 16)) & 0xF) - 8) & 0xFF; const int v3 = (((w >> (base + 24)) & 0xF) - 8) & 0xFF; - const uint n_col = n8_in_tile * 8u + r + parity * 4u; - const uint slab_idx = k4_in_chunk / (MMA_K >> 2u); - const uint k4_in_slab = k4_in_chunk % (MMA_K >> 2u); - Bsh_int8[slab_idx * B_SLAB_U32 + n_col * B_STRIDE_U32 + k4_in_slab] = + Bsh_int8[a] = uint(v0 | (v1 << 8) | (v2 << 16) | (v3 << 24)); } #else @@ -396,14 +430,12 @@ void main() { } #ifdef WEIGHT_INT4 [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { - const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; - const uint block_in_chunk = slot >> 3u; - const uint k4_blk = (chunkK_nxt >> 2u) + block_in_chunk / N8_PER_TILE; - const uint n8_blk = (tile_n_start >> 3u) + (block_in_chunk % N8_PER_TILE); + const uint a = gl_LocalInvocationID.x + si * WG_SIZE; + const BCoalIndex bidx1 = bcoal_index(a, chunkK_nxt, tile_n_start); #ifdef WEIGHT_BUFFER - temp_B[si] = t_packed_weight[(n8_blk * nblocks_x_A) + k4_blk]; + temp_B[si] = t_packed_weight[(bidx1.n8_blk * nblocks_x_A) + bidx1.k4_blk]; #else - temp_B[si] = texelFetch(t_packed_weight, ivec2(k4_blk, n8_blk), 0); + temp_B[si] = texelFetch(t_packed_weight, ivec2(bidx1.k4_blk, bidx1.n8_blk), 0); #endif } if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { @@ -467,23 +499,19 @@ void main() { } #ifdef WEIGHT_INT4 [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { - const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; - const uint block_in_chunk = slot >> 3u; - const uint col_in_block = slot & 7u; - const uint k4_in_chunk = block_in_chunk / N8_PER_TILE; - const uint n8_in_tile = block_in_chunk % N8_PER_TILE; - const uint r = col_in_block & 3u; - const uint parity = col_in_block >> 2u; - const int w = temp_B[si][r]; - const int base = int(4u * parity); + const uint a = gl_LocalInvocationID.x + si * WG_SIZE; + // r/parity depend only on n_col (a % B_SLAB_U32, mod 8) -- not on + // chunkK_base -- so this store site can pass a dummy 0u (unlike + // the fetch site above, a separate `if (has_next)` scope, which + // needs the real chunkK_nxt to compute n8_blk/k4_blk). + const BCoalIndex bidx1 = bcoal_index(a, 0u, tile_n_start); + const int w = temp_B[si][bidx1.r]; + const int base = int(4u * bidx1.parity); const int v0 = (((w >> (base + 0)) & 0xF) - 8) & 0xFF; const int v1 = (((w >> (base + 8)) & 0xF) - 8) & 0xFF; const int v2 = (((w >> (base + 16)) & 0xF) - 8) & 0xFF; const int v3 = (((w >> (base + 24)) & 0xF) - 8) & 0xFF; - const uint n_col = n8_in_tile * 8u + r + parity * 4u; - const uint slab_idx = k4_in_chunk / (MMA_K >> 2u); - const uint k4_in_slab = k4_in_chunk % (MMA_K >> 2u); - Bsh_int8[nxt_b + slab_idx * B_SLAB_U32 + n_col * B_STRIDE_U32 + k4_in_slab] = + Bsh_int8[nxt_b + a] = uint(v0 | (v1 << 8) | (v2 << 16) | (v3 << 24)); } if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp index c224925cd69..268d79e4919 100644 --- a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp @@ -89,7 +89,25 @@ constexpr CoopmatTileDims kDq8caQ4gswCoopmatDims = {64, 32, 32, 128, 2}; // variant (1-4); "tsweep_t..." is the original (production winner's own) // namespace. Unset/unrecognized = shipped dispatch, unchanged. The five // prefixes are mutually exclusive by construction (position 7 is 'd' vs 't'). -static const char* const kTsweepPrefixes[] = { +// Prefix lists are FAMILY-SCOPED. They used to be one shared kTsweepPrefixes +// array consulted by both selectors, which meant an 8da4w-only token supplied +// on ET_VK_Q4GSW_COOPMAT_VARIANT was "recognized", passed through, and then +// died at the later shader-lookup throw with a message that named the missing +// kernel rather than the actual mistake. Splitting them lets a wrong-family +// token be rejected here, by name. +// +// dbuf1-4 and the bare "tsweep_t" namespace exist for both families; the +// -tr/-zp/-zpn/-zpb variants are dq8ca-only (there is no q4gsw shader for +// any of them). +static const char* const kQ4gswTsweepPrefixes[] = { + "tsweep_dbuf1_t", + "tsweep_dbuf2_t", + "tsweep_dbuf3_t", + "tsweep_dbuf4_t", + "tsweep_t", +}; + +static const char* const kDq8caTsweepPrefixes[] = { "tsweep_dbuf1_t", "tsweep_dbuf2_t", "tsweep_dbuf3_t", @@ -97,8 +115,7 @@ static const char* const kTsweepPrefixes[] = { // "-tr": dbuf4's loop with the A-side global -> LDS staging swapped to // coopMatLoad/coopMatStore (ported from shmem_double_buf4-tr.comp). Stays // mutually exclusive with "tsweep_dbuf4_t" because position 12 is 't' vs - // '_', so prefix order here does not matter. dq8ca only -- there is no - // q4gsw "-tr" shader. + // '_', so prefix order here does not matter. "tsweep_dbuf4tr_t", // DIAGNOSTIC: same row-major layout as "-tr" but with scalar A staging, // used to bisect a dbuf4tr correctness failure. Delete with the shader. @@ -109,18 +126,59 @@ static const char* const kTsweepPrefixes[] = { "tsweep_dbuf4zp_t", // zp-hoisted + byte-parallel int4->int8 widening. "tsweep_dbuf4zpn_t", - // + templated B LDS skew (intervention A). + // + templated B LDS skew (intervention A). MEASURED-NEGATIVE: the + // bank-conflict skew is worth far more than the multiplies it costs. "tsweep_dbuf4zpb_t", + // zpn + ARITHMETIC nibble-parity select instead of a compared select + // (intervention E of dq8ca-prefill-stall-reduction). MEASURED-NEGATIVE: + // +2.87%; the variable shift is serially dependent on parity where the + // ternary's two constant shifts were not. + "tsweep_dbuf4zpx_t", + // zpn + loop-invariant staging index arithmetic hoisted out of the group + // loop (intervention F of dq8ca-prefill-stall-reduction). + "tsweep_dbuf4zpi_t", + // zpi + compile-time elision of the statically-true a_active guard + // (intervention G of dq8ca-prefill-stall-reduction). + "tsweep_dbuf4zpg_t", + // zpn + the a_active elision alone, no index hoist -- isolates G from F. + "tsweep_dbuf4zpk_t", + // (dq8ca-dequant-unpack-ablation and its 2026-08-26 follow-ups on + // xgpusw-debug08 -- abl_nodq/abl_nonib/abl_both/abl_nolds/abl_bconst/ + // abl_bcont/abl_breadc/str4/str6/str8/bcoal -- were measurement-only + // variants deleted once each attribution was recorded; see + // openspec/changes/dq8ca-dequant-unpack-ablation/results/. Two real + // findings from that investigation WERE promoted to the shipped default: + // see the B_STRIDE_U32 comment (the LDS skew removal) and the + // BCoalIndex/bcoal_index comment (the coalesced B-store rewrite) in + // linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4.glsl.) "tsweep_t", }; -static bool is_recognized_coopmat_variant_token(const std::string& v) { - for (const char* prefix : kTsweepPrefixes) { +// (The measurement-only ablation variants and their prefix list lived here +// during dq8ca-prefill-stall-reduction and were deleted together with their +// shaders once the attribution was recorded -- leaving numerically-wrong +// variants compiled into spv.cpp is a standing hazard. What they measured is in +// openspec/changes/dq8ca-prefill-stall-reduction/results/attribution.md.) + +// Returns the matched prefix's length, or npos. +template +static size_t match_prefix_len( + const std::string& v, + const char* const (&list)[N]) { + for (const char* prefix : list) { if (v.rfind(prefix, 0) == 0) { - return true; + return std::strlen(prefix); } } - return false; + return std::string::npos; +} + +static bool is_q4gsw_shippable_token(const std::string& v) { + return match_prefix_len(v, kQ4gswTsweepPrefixes) != std::string::npos; +} + +static bool is_dq8ca_shippable_token(const std::string& v) { + return match_prefix_len(v, kDq8caTsweepPrefixes) != std::string::npos; } static const std::string& q4gsw_coopmat_variant() { @@ -140,18 +198,41 @@ static const std::string& q4gsw_coopmat_variant() { return std::string("tsweep_dbuf4_t128x128k16g22s32"); } const std::string v(env); - if (is_recognized_coopmat_variant_token(v)) { + if (is_q4gsw_shippable_token(v)) { return v; } + // Reject a token that belongs to the OTHER shader family by name, rather + // than silently falling back (which hides the typo) or passing it through + // to die at the shader lookup (which reports the wrong problem). + VK_CHECK_COND( + !is_dq8ca_shippable_token(v), + "ET_VK_Q4GSW_COOPMAT_VARIANT was given '", + v, + "', which is a dq8ca (8da4w) shader-family token. There is no q4gsw " + "(4w) shader with that name. Use ET_VK_DQ8CA_COOPMAT_VARIANT instead."); return std::string("tsweep_dbuf4_t128x128k16g22s32"); }(); return variant; } static const std::string& dq8ca_coopmat_variant() { - // Default (no ET_VK_DQ8CA_COOPMAT_VARIANT set): tsweep_dbuf4_t64x32k32g12s64 - // -- same geometry as the shipped buffer-storage default, resolved through - // the tsweep_dbuf4 texture3d-capable shader. + // Default (no ET_VK_DQ8CA_COOPMAT_VARIANT set): + // tsweep_dbuf4_t128x128k64g81s64 (WG_TILE 128x128x64, SG_GRID 8x1, wave64) -- + // PROMOTED 2026-08-26 from the prior default tsweep_dbuf4_t64x32k32g12s64, + // together with the B_STRIDE_U32 skew removal in + // linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4.glsl (see that file's comment). + // Real, on-device measurement + // (openspec/changes/dq8ca-dequant-unpack-ablation/results/README.md, + // Addendum 8): this tile, WITH the stride fix, measures 36.27%/36.65%/36.97% + // efficiency of int8 peak on 1B/3B/8B prefill respectively (vs. 32.00-32.09% + // for the stride-fixed small tile alone) -- a further +11.8% to +13.2% + // relative speedup, consistent across all three model sizes, on top of the + // already-applied stride fix. Validated with the rigor the two incidents + // below establish as the actual bar for this kernel family: 10 consecutive + // `--correctness-only` passes, zero failures, for EVERY (model x storage) + // combination -- 1B/3B/8B x buffer/texture3d, 60 runs total, not just one + // storage mode or one pass. Decode and full e2e were explicitly out of + // scope (prefill-only, by request, to save time) -- not yet validated. // // REJECTED 2026-08-18: tsweep_dbuf4_t64x64k32g24s32 looked like a clean // win (correctness-first sweep pick, rep-confirmed faster on 1B/3B/8B @@ -192,13 +273,13 @@ static const std::string& dq8ca_coopmat_variant() { static const std::string variant = [] { const char* env = std::getenv("ET_VK_DQ8CA_COOPMAT_VARIANT"); if (!env) { - return std::string("tsweep_dbuf4_t64x32k32g12s64"); + return std::string("tsweep_dbuf4_t128x128k64g81s64"); } const std::string v(env); - if (is_recognized_coopmat_variant_token(v)) { + if (is_dq8ca_shippable_token(v)) { return v; } - return std::string("tsweep_dbuf4_t64x32k32g12s64"); + return std::string("tsweep_dbuf4_t128x128k64g81s64"); }(); return variant; } @@ -206,16 +287,13 @@ static const std::string& dq8ca_coopmat_variant() { // Parses "tsweep_txkgs" or // "tsweep_dbuf_txkgs" -> {M, N, K, SGX*SGY*sub, // SGY}. Returns fallback unchanged if the token matches none of -// kTsweepPrefixes. +// any family's prefix list. static CoopmatTileDims parse_tsweep_tile( const std::string& variant, const CoopmatTileDims& fallback) { - size_t t_pos = std::string::npos; - for (const char* prefix : kTsweepPrefixes) { - if (variant.rfind(prefix, 0) == 0) { - t_pos = std::strlen(prefix); - break; - } + size_t t_pos = match_prefix_len(variant, kDq8caTsweepPrefixes); + if (t_pos == std::string::npos) { + t_pos = match_prefix_len(variant, kQ4gswTsweepPrefixes); } if (t_pos == std::string::npos) { return fallback; From 249f8688714a19597835e25e14539ef2a2c6d7d9 Mon Sep 17 00:00:00 2001 From: Yanwen Xu Date: Fri, 28 Aug 2026 15:37:35 -0700 Subject: [PATCH 13/28] [ET-VK] Promote zpg (zp-hoist + coalesced B-store) as 8da4w coopmat default tsweep_dbuf4zpg_t128x64k32g42s32 combines zp-hoist, byte-parallel int4 unpack, static branch elision, and loop-invariant B-staging-index hoisting with the B_STRIDE_U32 skew removal and coalesced B-store the prior default already shipped. Validated on the sibling dq8ca-uvec4-redesign branch (cut from this branch at the same base commit) at 46.49-46.50% efficiency of int8 peak on 8B, vs. 36.97% for tsweep_dbuf4_t128x128k64g81s64 -- +20.5/20.9/21.2% relative on 8B/3B/1B. Re-verified on this branch's own build: 48/48 correctness cases pass (3 consecutive runs), and a real e2e run confirms the scored ETDump block dispatches this shader for 67.4% of leaf GPU time, 344.7 tok/s prefill on 8B (xgpusw-debug08, canonical main-fafb46ae9c0d driver, maxpin 980/5333/934). Removed the shaders and dispatch-table prefixes for the superseded single-intervention isolation variants (-tr/-trm/-trd, -zp/-zpn/-zpb/ -zpx, -zpi/-zpk) this promotion folds together or supersedes -- this branch ships only the validated default; the experimental siblings live on dq8ca-uvec4-redesign and dq8ca-arch-redesign. --- ...ar_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr.glsl | 715 ------------------ ...ar_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr.yaml | 188 ----- ...r_dq8ca_q4gsw_coopmat_tsweep_dbuf4trd.glsl | 682 ----------------- ...r_dq8ca_q4gsw_coopmat_tsweep_dbuf4trd.yaml | 188 ----- ...r_dq8ca_q4gsw_coopmat_tsweep_dbuf4trm.glsl | 707 ----------------- ...r_dq8ca_q4gsw_coopmat_tsweep_dbuf4trm.yaml | 188 ----- ...ar_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp.glsl | 687 ----------------- ...ar_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp.yaml | 187 ----- ...r_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpb.glsl | 714 ----------------- ...r_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpb.yaml | 42 - ..._dq8ca_q4gsw_coopmat_tsweep_dbuf4zpg.glsl} | 146 ++-- ...r_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpg.yaml | 80 ++ ...r_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpn.yaml | 187 ----- .../graph/ops/impl/QuantizedLinear.cpp | 79 +- 14 files changed, 220 insertions(+), 4570 deletions(-) delete mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr.glsl delete mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr.yaml delete mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trd.glsl delete mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trd.yaml delete mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trm.glsl delete mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trm.yaml delete mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp.glsl delete mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp.yaml delete mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpb.glsl delete mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpb.yaml rename backends/vulkan/runtime/graph/ops/glsl/{linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpn.glsl => linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpg.glsl} (85%) create mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpg.yaml delete mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpn.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr.glsl b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr.glsl deleted file mode 100644 index 102dea7bd07..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr.glsl +++ /dev/null @@ -1,715 +0,0 @@ -/* - * Copyright (c) Meta Platforms, Inc. and affiliates. - * All rights reserved. - * - * This source code is licensed under the BSD-style license found in the - * LICENSE file in the root directory of this source tree. - */ - -/* - * "-tr" (coopmat-staged A) variant of linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4. - * - * Ported from shmem_double_buf4-tr.comp on vk_cooperative_matrix_perf's - * gemm-ubm branch. That reference file's delta over shmem_double_buf4.comp is - * that the global -> LDS staging goes through COOPERATIVE MATRIX REGISTERS - * (coopMatLoad from global -> coopmat<> array -> coopMatStore into shared) - * instead of a hand-rolled per-thread uvec4 copy, and that B lands in LDS - * column-major. - * - * Only the A half of that idea is portable to this kernel: - * - * - B is ALREADY column-major in LDS here (Bsh_int8 is K-contiguous per - * output column, read back with gl_CooperativeMatrixLayoutColumnMajor), - * so the reference's "tr" property is not a delta for B at all. - * - B CANNOT be coopmat-staged: the weights are int4, each ivec4 holding - * 8 columns x 4 K-values that need the nibble-extract / -8 / sign-pack - * below. coopMatLoad cannot unpack nibbles, and a coopmat's per-lane - * layout is opaque so one cannot be built from unpacked registers. - * B staging is therefore left byte-identical to dbuf4. - * - A CAN be coopmat-staged, but only against a ROW-MAJOR packed int8 - * activation buffer. The stock 4h4w layout (kPackedInt8_4H4W, produced by - * quantize_and_pack_4h4w_with_group_sums.glsl) is NOT row-major: element - * [m4 * K4 + k4] is an ivec4 whose COMPONENT selects one of 4 rows, so as - * a uint array the index is m4*(4*K4) + k4*4 + r, which is not affine in - * the row index and cannot be addressed by any RowMajor/ColumnMajor - * coopMatLoad. (ColumnMajor is out on contiguity too: a uint packs 4 - * K-values, not 4 M-values.) - * - * So this shader binds t_packed_int8_input as a SCALAR int8_t array in the - * kPackedInt8_4W layout -- plain row-major int8, row stride K -- produced by quantize_and_pack_4w_with_group_sums.glsl. - * QuantizedLinear.cpp allocates that layout (and dispatches that packer) - * only when the active dq8ca variant is a "tsweep_dbuf4tr_t..." token AND - * the coopmat gate passes, so the tiled fallback never sees the wrong - * layout. Everything downstream of A staging -- LDS layout, int8 WMMA thread - * maps, group epilog, bias/store epilogue -- is unchanged from dbuf4. - * - * A staging (the actual -tr port): - * dbuf4: per-thread (m4, k4) ivec4 fetch; only A_ACTIVE_THREADS = - * (WG_TILE_M/4) * (WG_TILE_K/4) invocations participate, each - * scattering 4 rows into Ash_int8 with 4 scalar stores. - * dbuf4tr: per-SUBGROUP MMA_M x MMA_K tile fetch via coopMatLoad straight - * from global, then coopMatStore into the same Ash_int8 slot. - * The (WG_TILE_M/MMA_M) * (WG_TILE_K/MMA_K) tiles of a chunk are - * dealt round-robin across the NUM_SUBGROUPS subgroups. - * - * The loop structure is dbuf4's, unchanged: - * prologue: prefetch chunk 0 -> temp, store to slice 0 (no barrier) - * per iter: barrier -> prefetch(next) -> MMA(cur) -> store(next) - * and the nested `groups x chunks` loop with an unconditional group epilog is - * kept as-is (flattening it crashes the Xclipse PAL compiler at large - * spec-resolved trip counts -- see dbuf2's header). - * - * Selected at dispatch via - * ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4tr_txkgs<32|64> - * (QuantizedLinear.cpp), additive to the tsweep_dbuf4_t... and tsweep_t... - * namespaces. NOT the default -- unvalidated until it passes repeated - * test_llama_microbench --correctness-only runs (see dq8ca_coopmat_variant()'s - * comment on why a single pass is not proof). - * - * Performs: out[M,N] = dequant(int8_act) * dequant(int_w) (+ bias) - * via coopmat x coopmat -> coopmat on the matrix unit. - * - * Hard preconditions (in addition to dbuf4's): - * M % WG_TILE_M == 0, N % WG_TILE_N == 0, K % WG_TILE_K == 0, - * group_size % WG_TILE_K == 0, K % 4 == 0, - * WG_TILE_M % MMA_M == 0, WG_TILE_K % MMA_K == 0, - * t_packed_int8_input in kPackedInt8_4W (row-major) layout, - * device exposes coopmatx-> at 16x16x16. - */ - -#version 450 core - -#extension GL_KHR_cooperative_matrix : require -#extension GL_KHR_memory_scope_semantics : require -#extension GL_KHR_shader_subgroup_basic : enable -#extension GL_EXT_shader_explicit_arithmetic_types : require -#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require -// 8-bit SSBO access: A is bound as a scalar int8_t array so that the -// coopMatLoad below has a MATCHING component type. Loading a -// coopmat from a 32-bit int[] SSBO is what broke the first -// attempt (see header). -#extension GL_EXT_shader_8bit_storage : require -#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require -#extension GL_EXT_control_flow_attributes : enable - -#define PRECISION ${PRECISION} - -$if WEIGHT_NBITS == 4: - #define WEIGHT_INT4 - -$if HAS_BIAS: - #define HAS_BIAS - -$if WEIGHT_STORAGE == "buffer": - #define WEIGHT_BUFFER - -$if IO_STORAGE == "texture3d": - #define IO_TEXTURE - -layout(std430) buffer; - -#include "common.glslh" - -// Bindings — match add_linear_dqa_qw_node arg order: -// output(0), fp_input(1), packed_int8_input(2), int_input_sums(3 - unused), -// input_scales(4), input_zps(5), packed_weight(6), weight_sums(7), -// weight_scales(8), bias(9). -${layout_declare_tensor(B, "w", "t_output", "half", IO_STORAGE, is_scalar_array=True)} -// t_input is unread here -- the activations arrive already quantized in -// t_packed_int8_input -- but stays declared so the binding layout matches the -// dispatch site. It tracks IO_STORAGE so the two IO tensors stay consistent. -${layout_declare_tensor(B, "r", "t_input", "half", IO_STORAGE, is_scalar_array=False)} -// ROW-MAJOR (kPackedInt8_4W) packed activations, bound as a scalar int8_t -// array (row stride = K int8). Two things differ from dbuf4, which takes the -// 4h4w ivec4 block layout: -// 1. row-major, so a coopMatLoad can address it at all; -// 2. element type int8_t, MATCHING the coopmat component type. -// (2) is not cosmetic. Binding the same memory as int[] and loading a -// coopmat from it -- a type mismatch that demonstrably works for the -// Workgroup storage class, which is how the MMA loop reads Ash_int8 below -- -// silently produces wrong results from a StorageBuffer on this driver. -// The reference shmem_double_buf4-tr.comp sidesteps it the same way: its -// buffer_reference is declared `A_TYPE x[]`, i.e. int8_t for the int8 config. -// All A offsets/strides here are therefore in INT8 elements, not int. -${layout_declare_tensor(B, "r", "t_packed_int8_input", "int8", "buffer", is_scalar_array=True)} -${layout_declare_tensor(B, "r", "t_int8_input_sums", "int", "buffer", is_scalar_array=True)} -${layout_declare_tensor(B, "r", "t_int8_input_scales", "half", "texture3d")} -${layout_declare_tensor(B, "r", "t_int8_input_zps", "int8", "texture3d")} -${layout_declare_tensor(B, "r", "t_packed_weight", "int", WEIGHT_STORAGE, is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_weight_sums", "int", "buffer", is_scalar_array=True)} -${layout_declare_tensor(B, "r", "t_weight_scales", "half", "buffer", is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_bias", "half", "buffer", is_scalar_array=True)} - -${layout_declare_ubo(B, "ivec4", "output_sizes")} -${layout_declare_ubo(B, "ivec4", "input_sizes")} - -layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; - -${layout_declare_spec_const(C, "int", "apply_bias", "0")} -// INT4 only; inert (0) for INT8 so the dispatcher's spec list lines up. -${layout_declare_spec_const(C, "int", "K4_per_group", "0")} -${layout_declare_spec_const(C, "int", "num_groups_arg", "0")} -${layout_declare_spec_const(C, "int", "out_N_arg", "0")} - -// Tile geometry -const uint MMA_M = ${MMA_M}; -const uint MMA_N = ${MMA_N}; -const uint MMA_K = ${MMA_K}; - -const uint WG_TILE_M = ${WG_TILE_M}; -const uint WG_TILE_N = ${WG_TILE_N}; -const uint WG_TILE_K = ${WG_TILE_K}; - -const uint SG_GRID_X = ${SG_GRID_X}; -const uint SG_GRID_Y = ${SG_GRID_Y}; -const uint SUBGROUP_SIZE = ${SUBGROUP_SIZE}; -const uint NUM_SUBGROUPS = SG_GRID_X * SG_GRID_Y; -const uint WG_SIZE = NUM_SUBGROUPS * SUBGROUP_SIZE; - -const uint SG_TILE_M = WG_TILE_M / SG_GRID_Y; -const uint SG_TILE_N = WG_TILE_N / SG_GRID_X; -const uint MMAS_PER_SG_M = SG_TILE_M / MMA_M; -const uint MMAS_PER_SG_N = SG_TILE_N / MMA_N; - -const uint A_SLAB_INT8 = WG_TILE_M * MMA_K; -const uint B_USEFUL_U32 = MMA_K / 4u; -const uint B_STRIDE_U32 = B_USEFUL_U32 + 1u; // +1 skew -const uint B_SLAB_U32 = WG_TILE_N * B_STRIDE_U32; -const uint NUM_K_SLABS = WG_TILE_K / MMA_K; - -const uint A_SLAB_U32 = A_SLAB_INT8 / 4u; -const uint A_STRIDE_U32 = MMA_K / 4u; - -// One ping-pong slice covers all K-slabs of one chunk. -const uint ASH_SLICE_U32 = NUM_K_SLABS * A_SLAB_U32; -const uint BSH_SLICE_U32 = NUM_K_SLABS * B_SLAB_U32; - -// Double-buffered MMA operand staging. -shared uint Ash_int8[2u * ASH_SLICE_U32]; -shared uint Bsh_int8[2u * BSH_SLICE_U32]; - -// Per-WG-tile-row activation params (loaded ONCE at WG start; constant -// across groups). -shared int izp_sh[WG_TILE_M]; // int32 (cast from int8 source) for broadcast -shared float ifs_sh[WG_TILE_M]; // float32 (cast from fp16 source) for broadcast - -// Per-(group, output-channel) weight params, ping-ponged by group parity. -// (For per-channel INT8 only slice 0 is ever used.) -shared int wsum_sh[2u * WG_TILE_N]; -shared float wsc_sh[2u * WG_TILE_N]; - -#ifdef HAS_BIAS -shared float bias_sh[WG_TILE_N]; -#endif - -#ifdef IO_TEXTURE -// Result staging for the imageStore epilogue, mirroring the fp16 kernel: -// SG_GRID_Y bands of MMA_M rows, each WG_TILE_N wide, row-major. A full -// WG_TILE_M x WG_TILE_N buffer would cost SG_GRID_Y/MMAS_PER_SG_M x more LDS -// and wreck occupancy. float16_t-typed because coopMatStore needs it. -const uint CSH_ROWS = SG_GRID_Y * MMA_M; -shared float16_t Csh_out[CSH_ROWS * WG_TILE_N]; -#endif - -// Running fp32 accumulator (across all groups). -coopmat - result[MMAS_PER_SG_M][MMAS_PER_SG_N]; - -// Per-group int32 MMA accumulator. -coopmat - accum_int32[MMAS_PER_SG_M][MMAS_PER_SG_N]; - -void main() { - const uvec2 tileID = uvec2(gl_WorkGroupID.xy); - const uvec2 warpInTile = uvec2( - gl_SubgroupID % SG_GRID_X, - gl_SubgroupID / SG_GRID_X); - - const uint K = uint(input_sizes.x); - const uint N = uint(output_sizes.x); - const uint N4 = (N + 3u) / 4u; - const uint nblocks_x_A = (K + 3u) >> 2u; - // A row stride in INT8 elements (the binding's element type). Derived from - // nblocks_x_A rather than K directly so it matches the packer's - // `m_row * K4 + k4` addressing exactly; K %% 4 == 0 makes them equal. - const uint a_row_stride_i8 = nblocks_x_A * 4u; - -#ifdef WEIGHT_INT4 - const uint num_groups = uint(num_groups_arg); - const uint CHUNKS_PER_GROUP = uint(K4_per_group) * 4u / WG_TILE_K; -#else - // Per-channel: a single quant "group" spanning all of K. The nested - // groups x chunks loop below collapses to a flat chunk loop, the wsum/wsc - // ping-pong never crosses a boundary, and the epilog runs exactly once. - const uint num_groups = 1u; - const uint CHUNKS_PER_GROUP = uint(num_groups_arg); -#endif - const uint num_chunks = num_groups * CHUNKS_PER_GROUP; - - const uint tile_m_start = WG_TILE_M * tileID.y; - const uint tile_n_start = WG_TILE_N * tileID.x; - - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - result[i][j] = coopmat(0.0); - accum_int32[i][j] = coopmat(0); - } - } - - const uint K_BLOCKS_PER_CHUNK = WG_TILE_K >> 2u; - - // --- A staging tile map: one MMA_M x MMA_K coopmat tile per subgroup per - // slot. A chunk holds A_TILES_M x A_TILES_K such tiles; they are dealt - // round-robin across the NUM_SUBGROUPS subgroups, so every subgroup - // participates (dbuf4's per-thread map leaves WG_SIZE - - // A_ACTIVE_THREADS invocations idle whenever the tile is small). - // A_TILES_PER_SG rounds up, so the last slot may be partially used -- - // the `t < NUM_A_TILES` guard below is subgroup-uniform (t depends only - // on gl_SubgroupID), which is what coopmat ops require. - const uint A_TILES_M = WG_TILE_M / MMA_M; - const uint A_TILES_K = WG_TILE_K / MMA_K; // == NUM_K_SLABS - const uint NUM_A_TILES = A_TILES_M * A_TILES_K; - const uint A_TILES_PER_SG = (NUM_A_TILES + NUM_SUBGROUPS - 1u) / NUM_SUBGROUPS; - -#ifdef WEIGHT_INT4 - // --- B staging thread map: (block, col) slots; each slot extracts one - // ColumnMajor LDS uint (4 K-contiguous sign-extended int8) --- - const uint B_TOTAL_SLOTS = K_BLOCKS_PER_CHUNK * WG_TILE_N; - const uint B_SLOTS_PER_THREAD = B_TOTAL_SLOTS / WG_SIZE; - const uint N8_PER_TILE = WG_TILE_N >> 3u; -#else - // --- B staging thread map: one (k4, n4) ivec4 block per active thread --- - const uint B_FETCH_SLOTS = K_BLOCKS_PER_CHUNK * (WG_TILE_N >> 2u); - const uint N4_PER_TILE = WG_TILE_N >> 2u; - const uint b_k4_in_chunk = gl_LocalInvocationID.x / N4_PER_TILE; - const uint b_n_uint_col = gl_LocalInvocationID.x % N4_PER_TILE; - const bool b_active = gl_LocalInvocationID.x < B_FETCH_SLOTS; -#endif - - // Prefetch temp registers. temp_A is a coopmat array (the -tr change); - // indices into it are [[unroll]]-resolved compile-time constants, never - // dynamic -- dynamic indexing of a coopmat array is exactly the construct - // the Xclipse/AMD-PAL compiler has miscompiled before. - coopmat - temp_A[A_TILES_PER_SG]; -#ifdef WEIGHT_INT4 - ivec4 temp_B[B_SLOTS_PER_THREAD]; - int temp_wsum; - float temp_wsc; -#else - ivec4 temp_B; -#endif - - // ========================================================= - // PROLOGUE - // ========================================================= - if (gl_LocalInvocationID.x < (WG_TILE_M >> 2u)) { - const uint m4 = (tile_m_start >> 2u) + gl_LocalInvocationID.x; - const vec4 sc = vec4(texelFetch(t_int8_input_scales, ivec3(m4, 0, 0), 0)); - const ivec4 zp = texelFetch(t_int8_input_zps, ivec3(m4, 0, 0), 0); - const uint base = gl_LocalInvocationID.x * 4u; - ifs_sh[base + 0u] = sc.x; ifs_sh[base + 1u] = sc.y; - ifs_sh[base + 2u] = sc.z; ifs_sh[base + 3u] = sc.w; - izp_sh[base + 0u] = zp.x; izp_sh[base + 1u] = zp.y; - izp_sh[base + 2u] = zp.z; izp_sh[base + 3u] = zp.w; - } - // Group 0 weight sums/scales -> slice 0. - if (gl_LocalInvocationID.x < WG_TILE_N) { - const uint n_idx = tile_n_start + gl_LocalInvocationID.x; - f16vec4 sv = t_weight_scales[n_idx >> 2u]; - wsc_sh[gl_LocalInvocationID.x] = float(sv[n_idx & 3u]); - wsum_sh[gl_LocalInvocationID.x] = t_weight_sums[n_idx]; - } - memoryBarrierShared(); - barrier(); - - // izp/ifs are per-row activation params, constant across K groups — - // broadcast them into coopmats ONCE; the group epilog reuses them every - // group (they depend only on the row block i, not on the group or j). - coopmat - izp_bcast[MMAS_PER_SG_M]; - coopmat - ifs_bcast[MMAS_PER_SG_M]; - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - const uint local_m_base = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); - coopMatLoad( - izp_bcast[i], izp_sh, - local_m_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutColumnMajor); - coopMatLoad( - ifs_bcast[i], ifs_sh, - local_m_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutColumnMajor); - } - - // dbuf4: prefetch chunk 0 into temp registers, THEN store to slice 0 (no - // barrier here -- the main loop's first iteration barriers before - // reading slice 0). - [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { - const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; - if (t < NUM_A_TILES) { - const uint tm = t / A_TILES_K; - const uint tk = t % A_TILES_K; - // Offset and stride are in ARRAY ELEMENTS, which for this binding are - // int8 -- i.e. the natural row-major coordinates. - coopMatLoad( - temp_A[s], t_packed_int8_input, - (tile_m_start + tm * MMA_M) * a_row_stride_i8 + tk * MMA_K, - a_row_stride_i8, - gl_CooperativeMatrixLayoutRowMajor); - } - } -#ifdef WEIGHT_INT4 - [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { - const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; - const uint block_in_chunk = slot >> 3u; - const uint k4_blk = block_in_chunk / N8_PER_TILE; - const uint n8_blk = (tile_n_start >> 3u) + (block_in_chunk % N8_PER_TILE); -#ifdef WEIGHT_BUFFER - temp_B[si] = t_packed_weight[(n8_blk * nblocks_x_A) + k4_blk]; -#else - temp_B[si] = texelFetch(t_packed_weight, ivec2(k4_blk, n8_blk), 0); -#endif - } -#else - if (b_active) { - const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; -#ifdef WEIGHT_BUFFER - temp_B = t_packed_weight[(b_k4_in_chunk * N4) + block_x_w]; -#else - temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, b_k4_in_chunk), 0); -#endif - } -#endif - { - // store chunk 0 -> slice 0 - [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { - const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; - if (t < NUM_A_TILES) { - const uint tm = t / A_TILES_K; - const uint tk = t % A_TILES_K; - coopMatStore( - temp_A[s], Ash_int8, - tk * A_SLAB_U32 + (tm * MMA_M) * A_STRIDE_U32, - A_STRIDE_U32, - gl_CooperativeMatrixLayoutRowMajor); - } - } -#ifdef WEIGHT_INT4 - [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { - const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; - const uint block_in_chunk = slot >> 3u; - const uint col_in_block = slot & 7u; - const uint k4_in_chunk = block_in_chunk / N8_PER_TILE; - const uint n8_in_tile = block_in_chunk % N8_PER_TILE; - const uint r = col_in_block & 3u; - const uint parity = col_in_block >> 2u; - const int w = temp_B[si][r]; - const int base = int(4u * parity); - const int v0 = (((w >> (base + 0)) & 0xF) - 8) & 0xFF; - const int v1 = (((w >> (base + 8)) & 0xF) - 8) & 0xFF; - const int v2 = (((w >> (base + 16)) & 0xF) - 8) & 0xFF; - const int v3 = (((w >> (base + 24)) & 0xF) - 8) & 0xFF; - const uint n_col = n8_in_tile * 8u + r + parity * 4u; - const uint slab_idx = k4_in_chunk / (MMA_K >> 2u); - const uint k4_in_slab = k4_in_chunk % (MMA_K >> 2u); - Bsh_int8[slab_idx * B_SLAB_U32 + n_col * B_STRIDE_U32 + k4_in_slab] = - uint(v0 | (v1 << 8) | (v2 << 16) | (v3 << 24)); - } -#else - if (b_active) { - const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); - const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); - const uint n_col_base = b_n_uint_col * 4u; - [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { - Bsh_int8[slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab] = - uint(temp_B[n_in_blk]); - } - } -#endif - } - - // ========================================================= - // MAIN LOOP (dbuf4) — nested groups x chunks (kept nested; flattening it - // with a conditional coopmat epilog crashes the Xclipse PAL compiler at - // large spec-resolved trip counts). One barrier per chunk. Chunk - // iteration (global index `chunk`): - // 1. barrier — A/B slice (chunk%2) fully written; on the first chunk - // of group g, wsum/wsc slice (g%2) is too. - // 2. prefetch — chunk+1 (A blocks, B blocks) into temp; when chunk+1 - // starts a new group, also its wsum/wsc element. Skipped - // entirely on the final chunk. - // 3. int8 MMA — on slice (chunk%2) into accum_int32. - // 4. store — temp -> A/B slice ((chunk+1)%2), unpacking the weight; - // on a group boundary, wsum/wsc -> slice ((g+1)%2). - // The group epilog runs unconditionally at the tail of each group. - // ========================================================= - uint chunk = 0; - for (uint group_i = 0; group_i < num_groups; ++group_i) { - for (uint inner = 0; inner < CHUNKS_PER_GROUP; ++inner, ++chunk) { - const bool has_next = chunk + 1u < num_chunks; - const bool group_crossing = has_next && (inner + 1u == CHUNKS_PER_GROUP); - const uint cur_a = (chunk % 2u) * ASH_SLICE_U32; - const uint cur_b = (chunk % 2u) * BSH_SLICE_U32; - const uint nxt_a = ((chunk + 1u) % 2u) * ASH_SLICE_U32; - const uint nxt_b = ((chunk + 1u) % 2u) * BSH_SLICE_U32; - - barrier(); - - // --- 2. prefetch chunk+1 -> temp --- - if (has_next) { - const uint chunkK_nxt = (chunk + 1u) * WG_TILE_K; - [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { - const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; - if (t < NUM_A_TILES) { - const uint tm = t / A_TILES_K; - const uint tk = t % A_TILES_K; - coopMatLoad( - temp_A[s], t_packed_int8_input, - (tile_m_start + tm * MMA_M) * a_row_stride_i8 + chunkK_nxt + - tk * MMA_K, - a_row_stride_i8, - gl_CooperativeMatrixLayoutRowMajor); - } - } -#ifdef WEIGHT_INT4 - [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { - const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; - const uint block_in_chunk = slot >> 3u; - const uint k4_blk = (chunkK_nxt >> 2u) + block_in_chunk / N8_PER_TILE; - const uint n8_blk = (tile_n_start >> 3u) + (block_in_chunk % N8_PER_TILE); -#ifdef WEIGHT_BUFFER - temp_B[si] = t_packed_weight[(n8_blk * nblocks_x_A) + k4_blk]; -#else - temp_B[si] = texelFetch(t_packed_weight, ivec2(k4_blk, n8_blk), 0); -#endif - } - if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { - const uint n_idx = tile_n_start + gl_LocalInvocationID.x; - f16vec4 sv = t_weight_scales[(group_i + 1u) * N4 + (n_idx >> 2u)]; - temp_wsc = float(sv[n_idx & 3u]); - temp_wsum = t_weight_sums[(group_i + 1u) * N + n_idx]; - } -#else - if (b_active) { - const uint block_y_w = (chunkK_nxt >> 2u) + b_k4_in_chunk; - const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; -#ifdef WEIGHT_BUFFER - temp_B = t_packed_weight[(block_y_w * N4) + block_x_w]; -#else - temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, block_y_w), 0); -#endif - } -#endif - } - - // --- 3. int8 MMA on the cur slice --- - [[unroll]] for (uint k = 0; k < NUM_K_SLABS; ++k) { - const uint slab_a_base_u32 = cur_a + k * A_SLAB_U32; - const uint slab_b_base_u32 = cur_b + k * B_SLAB_U32; - - coopmat matA[MMAS_PER_SG_M]; - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - const uint row_a = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); - coopMatLoad( - matA[i], Ash_int8, - slab_a_base_u32 + row_a * A_STRIDE_U32, - A_STRIDE_U32, - gl_CooperativeMatrixLayoutRowMajor); - } - - coopmat matB; - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint col_b = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopMatLoad( - matB, Bsh_int8, - slab_b_base_u32 + col_b * B_STRIDE_U32, - B_STRIDE_U32, - gl_CooperativeMatrixLayoutColumnMajor); - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - accum_int32[i][j] = coopMatMulAdd(matA[i], matB, accum_int32[i][j]); - } - } - } - - // --- 4. store temp (chunk+1) -> nxt slice --- - if (has_next) { - [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { - const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; - if (t < NUM_A_TILES) { - const uint tm = t / A_TILES_K; - const uint tk = t % A_TILES_K; - coopMatStore( - temp_A[s], Ash_int8, - nxt_a + tk * A_SLAB_U32 + (tm * MMA_M) * A_STRIDE_U32, - A_STRIDE_U32, - gl_CooperativeMatrixLayoutRowMajor); - } - } -#ifdef WEIGHT_INT4 - [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { - const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; - const uint block_in_chunk = slot >> 3u; - const uint col_in_block = slot & 7u; - const uint k4_in_chunk = block_in_chunk / N8_PER_TILE; - const uint n8_in_tile = block_in_chunk % N8_PER_TILE; - const uint r = col_in_block & 3u; - const uint parity = col_in_block >> 2u; - const int w = temp_B[si][r]; - const int base = int(4u * parity); - const int v0 = (((w >> (base + 0)) & 0xF) - 8) & 0xFF; - const int v1 = (((w >> (base + 8)) & 0xF) - 8) & 0xFF; - const int v2 = (((w >> (base + 16)) & 0xF) - 8) & 0xFF; - const int v3 = (((w >> (base + 24)) & 0xF) - 8) & 0xFF; - const uint n_col = n8_in_tile * 8u + r + parity * 4u; - const uint slab_idx = k4_in_chunk / (MMA_K >> 2u); - const uint k4_in_slab = k4_in_chunk % (MMA_K >> 2u); - Bsh_int8[nxt_b + slab_idx * B_SLAB_U32 + n_col * B_STRIDE_U32 + k4_in_slab] = - uint(v0 | (v1 << 8) | (v2 << 16) | (v3 << 24)); - } - if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { - const uint wbase_nxt = ((group_i + 1u) % 2u) * WG_TILE_N; - wsum_sh[wbase_nxt + gl_LocalInvocationID.x] = temp_wsum; - wsc_sh[wbase_nxt + gl_LocalInvocationID.x] = temp_wsc; - } -#else - if (b_active) { - const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); - const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); - const uint n_col_base = b_n_uint_col * 4u; - [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { - Bsh_int8[nxt_b + slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab] = - uint(temp_B[n_in_blk]); - } - } -#endif - } - } // chunks - - // --- Group epilog: dequant accum_int32 -> result, reset accum --- - { - const uint wbase = (group_i % 2u) * WG_TILE_N; - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - - coopmat wsum_bcast; - coopMatLoad( - wsum_bcast, wsum_sh, - wbase + local_n_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutRowMajor); - - coopmat wsc_bcast; - coopMatLoad( - wsc_bcast, wsc_sh, - wbase + local_n_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutRowMajor); - - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - coopmat adjusted = - accum_int32[i][j] - izp_bcast[i] * wsum_bcast; - coopmat adjusted_fp = - coopmat(adjusted); - coopmat scales_outer = - ifs_bcast[i] * wsc_bcast; - result[i][j] += adjusted_fp * scales_outer; - accum_int32[i][j] = coopmat(0); - } - } - } - } // groups - - // --- Bias (optional) --- -#ifdef HAS_BIAS - if (apply_bias > 0) { - for (uint t = gl_LocalInvocationID.x; t < WG_TILE_N; t += WG_SIZE) { - bias_sh[t] = float(t_bias[tile_n_start + t]); - } - memoryBarrierShared(); - barrier(); - } -#endif - - // --- Store result tile --- - // N for the store address math MUST come from the spec constant, not the - // sizes UBO (see out_N_arg above). -#ifdef IO_TEXTURE - // Epilogue iteration i drains accumulator row-block i from EVERY subgroup - // into Csh_out at once, so the SG_GRID_Y bands it holds are disjoint global - // row ranges; the whole workgroup then imageStores them. lr / MMA_M is the - // writing subgroup's warpInTile.y, so the global row reproduces the buffer - // path's gi exactly. - // - // PORTABILITY NOTE: the barrier() in the loop body keeps this loop rolled - // despite [[unroll]], so result[i][j] IS dynamically indexed. Coopmat arrays - // are opaque per-lane storage and dynamic indexing is exactly the construct - // the Xclipse/AMD-PAL compiler has broken before -- check this first if the - // texture variants miscompile on M51. - const uint CSH_TEXELS_PER_ROW = WG_TILE_N / 4u; - const uint CSH_TEXELS = CSH_ROWS * CSH_TEXELS_PER_ROW; - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - // Guards Csh_out against the previous iteration's readers. Inert on i == 0 - // but must stay unconditional to remain workgroup-uniform. - barrier(); - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { -#ifdef HAS_BIAS - if (apply_bias > 0) { - const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopmat bias_tile; - coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); - result[i][j] += bias_tile; - } -#endif - coopmat out_tile = - coopmat(result[i][j]); - coopMatStore( - out_tile, Csh_out, - warpInTile.y * MMA_M * WG_TILE_N + - MMA_N * (MMAS_PER_SG_N * warpInTile.x + j), - WG_TILE_N, - gl_CooperativeMatrixLayoutRowMajor); - } - memoryBarrierShared(); - barrier(); - - for (uint t = gl_LocalInvocationID.x; t < CSH_TEXELS; t += WG_SIZE) { - const uint lr = t / CSH_TEXELS_PER_ROW; - const uint lc4 = t % CSH_TEXELS_PER_ROW; - const uint m = - tile_m_start + (lr / MMA_M) * SG_TILE_M + i * MMA_M + (lr % MMA_M); - const uint base = lr * WG_TILE_N + lc4 * 4u; - imageStore( - t_output, - ivec3(tile_n_start / 4u + lc4, m, 0), - vec4( - float(Csh_out[base]), - float(Csh_out[base + 1u]), - float(Csh_out[base + 2u]), - float(Csh_out[base + 3u]))); - } - } -#else - const uint N_out = uint(out_N_arg); - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint gi = tile_m_start + MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); - const uint gj = tile_n_start + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - -#ifdef HAS_BIAS - if (apply_bias > 0) { - const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopmat bias_tile; - coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); - result[i][j] += bias_tile; - } -#endif - - coopmat out_tile = - coopmat(result[i][j]); - coopMatStore( - out_tile, t_output, - gi * N_out + gj, N_out, - gl_CooperativeMatrixLayoutRowMajor); - } - } -#endif // IO_TEXTURE -} diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr.yaml b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr.yaml deleted file mode 100644 index f3f6c1d7773..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr.yaml +++ /dev/null @@ -1,188 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -# "-tr" (coopmat-staged A) variants of the int8 dq8ca_q4gsw coopmat kernel, -# ported from shmem_double_buf4-tr.comp (vk_cooperative_matrix_perf, -# gemm-ubm). linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr.glsl is a fork of -# linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4.glsl with ONLY the A-side global -> -# LDS staging swapped from a per-thread ivec4 copy to -# coopMatLoad(global) -> coopmat<> -> coopMatStore(LDS). B staging is -# byte-identical (int4 nibble unpack cannot be coopmat-staged) and B was -# already column-major in LDS, so it is not a delta. -# -# These variants require t_packed_int8_input in the ROW-MAJOR -# kPackedInt8_4W layout (produced by quantize_and_pack_4w_with_group_sums), -# NOT the stock 4h4w block layout. QuantizedLinear.cpp switches the tensor -# layout and the packer node together with the variant token, and only when -# the coopmat gate passes -- see dq8ca_wants_rowmajor_int8_input(). -# -# Selected via -# ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4tr_txkgs. -# NOT the shipped default: unvalidated until it passes repeated -# test_llama_microbench --correctness-only runs (a single pass is not proof -- -# see dq8ca_coopmat_variant()'s comment). -# -# Tile-geometry preconditions beyond dbuf4's: WG_TILE_M % MMA_M == 0 and -# WG_TILE_K % MMA_K == 0 (both hold for every entry below). The seed set is -# deliberately small -- widen it from a sweep once correctness is established. - -linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr: - parameter_names_with_default_values: - PRECISION: highp - HAS_BIAS: false - IO_STORAGE: buffer - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - MMA_M: 16 - MMA_N: 16 - MMA_K: 16 - WG_TILE_M: 64 - WG_TILE_N: 32 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - shader_variants: - # production 8da4w tile (dbuf2/dbuf4 default) -- the A/B anchor - # A tiles/chunk = 8, subgroups = 2 -> A_TILES_PER_SG = 4 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr_t64x32k32g12s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 32 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr_t64x32k32g12s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 32 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr_t64x32k32g12s64_buffer_buffer_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: buffer - WG_TILE_M: 64 - WG_TILE_N: 32 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - # q4gsw production geometry; 8 A-tiles over 4 subgroups - # A tiles/chunk = 8, subgroups = 4 -> A_TILES_PER_SG = 2 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr_t128x128k16g22s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 2 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr_t128x128k16g22s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 2 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - # 1 A-tile per subgroup (A_TILES_PER_SG == 1, no leftover slot) - # A tiles/chunk = 8, subgroups = 8 -> A_TILES_PER_SG = 1 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr_t64x64k32g24s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr_t64x64k32g24s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - # 16 A-tiles over 8 subgroups - # A tiles/chunk = 16, subgroups = 8 -> A_TILES_PER_SG = 2 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr_t128x64k32g24s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr_t128x64k32g24s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - # Reference-matching tile: shmem_double_buf4-tr.comp's own geometry - # (TILE 128x128x16, WORKGROUP_WIDTH_IN_SUBGROUPS=4 x DBUF4_GRID_HEIGHT=2, - # wave32) => 8 subgroups / 256 threads, MMAS_PER_SG 4x2. Added so our - # kernel can be A/B'd against the teammate's PAL capture at the same - # tile and the same 2048x1024x4096 shape (8b wk_wv prefill). - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr_t128x128k16g42s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr_t128x128k16g42s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - # G2/Pavan's documented INT8 uBM config: TILE 128x128x64 with a 4x2 - # subgroup grid at wave32 (workgroupSize=256). Added to compare our kernel - # against their reference dbuf4 number (~1973us) at the same geometry and - # the same shape (M2048 K4096 N1024 == 8b wk_wv prefill). - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr_t128x128k64g42s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr_t128x128k64g42s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trd.glsl b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trd.glsl deleted file mode 100644 index 2ecc26ce0f9..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trd.glsl +++ /dev/null @@ -1,682 +0,0 @@ -/* - * Copyright (c) Meta Platforms, Inc. and affiliates. - * All rights reserved. - * - * This source code is licensed under the BSD-style license found in the - * LICENSE file in the root directory of this source tree. - */ - -/* - * DIAGNOSTIC BISECT variant #2 of linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr. - * - * A does not go through LDS at all: each subgroup coopMatLoads the A tiles it - * needs straight from the row-major global buffer inside the MMA loop. This - * isolates coopMatLoad-from-StorageBuffer-int[] from - * coopMatStore-into-Workgroup-uint[]: dbuf4trm (manual staging) already proved - * the packer and layout correct, so if this variant passes the load is fine - * and the store is the culprit, and if it fails the load is. - * - * Ash_int8 is left allocated but unused (keeps the diff minimal; costs LDS). - * - * Original dbuf4tr header follows. - * - * "-tr" (coopmat-staged A) variant of linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4. - * - * Ported from shmem_double_buf4-tr.comp on vk_cooperative_matrix_perf's - * gemm-ubm branch. That reference file's delta over shmem_double_buf4.comp is - * that the global -> LDS staging goes through COOPERATIVE MATRIX REGISTERS - * (coopMatLoad from global -> coopmat<> array -> coopMatStore into shared) - * instead of a hand-rolled per-thread uvec4 copy, and that B lands in LDS - * column-major. - * - * Only the A half of that idea is portable to this kernel: - * - * - B is ALREADY column-major in LDS here (Bsh_int8 is K-contiguous per - * output column, read back with gl_CooperativeMatrixLayoutColumnMajor), - * so the reference's "tr" property is not a delta for B at all. - * - B CANNOT be coopmat-staged: the weights are int4, each ivec4 holding - * 8 columns x 4 K-values that need the nibble-extract / -8 / sign-pack - * below. coopMatLoad cannot unpack nibbles, and a coopmat's per-lane - * layout is opaque so one cannot be built from unpacked registers. - * B staging is therefore left byte-identical to dbuf4. - * - A CAN be coopmat-staged, but only against a ROW-MAJOR packed int8 - * activation buffer. The stock 4h4w layout (kPackedInt8_4H4W, produced by - * quantize_and_pack_4h4w_with_group_sums.glsl) is NOT row-major: element - * [m4 * K4 + k4] is an ivec4 whose COMPONENT selects one of 4 rows, so as - * a uint array the index is m4*(4*K4) + k4*4 + r, which is not affine in - * the row index and cannot be addressed by any RowMajor/ColumnMajor - * coopMatLoad. (ColumnMajor is out on contiguity too: a uint packs 4 - * K-values, not 4 M-values.) - * - * So this shader binds t_packed_int8_input as a SCALAR int array in the - * kPackedInt8_4W layout -- plain row-major int8, 4 K-values per int32, - * row stride K4 -- produced by quantize_and_pack_4w_with_group_sums.glsl. - * QuantizedLinear.cpp allocates that layout (and dispatches that packer) - * only when the active dq8ca variant is a "tsweep_dbuf4tr_t..." token AND - * the coopmat gate passes, so the tiled fallback never sees the wrong - * layout. Everything downstream of A staging -- LDS layout, int8 WMMA thread - * maps, group epilog, bias/store epilogue -- is unchanged from dbuf4. - * - * A staging (the actual -tr port): - * dbuf4: per-thread (m4, k4) ivec4 fetch; only A_ACTIVE_THREADS = - * (WG_TILE_M/4) * (WG_TILE_K/4) invocations participate, each - * scattering 4 rows into Ash_int8 with 4 scalar stores. - * dbuf4tr: per-SUBGROUP MMA_M x MMA_K tile fetch via coopMatLoad straight - * from global, then coopMatStore into the same Ash_int8 slot. - * The (WG_TILE_M/MMA_M) * (WG_TILE_K/MMA_K) tiles of a chunk are - * dealt round-robin across the NUM_SUBGROUPS subgroups. - * - * The loop structure is dbuf4's, unchanged: - * prologue: prefetch chunk 0 -> temp, store to slice 0 (no barrier) - * per iter: barrier -> prefetch(next) -> MMA(cur) -> store(next) - * and the nested `groups x chunks` loop with an unconditional group epilog is - * kept as-is (flattening it crashes the Xclipse PAL compiler at large - * spec-resolved trip counts -- see dbuf2's header). - * - * Selected at dispatch via - * ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4tr_txkgs<32|64> - * (QuantizedLinear.cpp), additive to the tsweep_dbuf4_t... and tsweep_t... - * namespaces. NOT the default -- unvalidated until it passes repeated - * test_llama_microbench --correctness-only runs (see dq8ca_coopmat_variant()'s - * comment on why a single pass is not proof). - * - * Performs: out[M,N] = dequant(int8_act) * dequant(int_w) (+ bias) - * via coopmat x coopmat -> coopmat on the matrix unit. - * - * Hard preconditions (in addition to dbuf4's): - * M % WG_TILE_M == 0, N % WG_TILE_N == 0, K % WG_TILE_K == 0, - * group_size % WG_TILE_K == 0, K % 4 == 0, - * WG_TILE_M % MMA_M == 0, WG_TILE_K % MMA_K == 0, - * t_packed_int8_input in kPackedInt8_4W (row-major) layout, - * device exposes coopmatx-> at 16x16x16. - */ - -#version 450 core - -#extension GL_KHR_cooperative_matrix : require -#extension GL_KHR_memory_scope_semantics : require -#extension GL_KHR_shader_subgroup_basic : enable -#extension GL_EXT_shader_explicit_arithmetic_types : require -#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require -// 8-bit SSBO access: A is bound as a scalar int8_t array so that the -// coopMatLoad below has a MATCHING component type. Loading a -// coopmat from a 32-bit int[] SSBO is what broke the first -// attempt (see header). -#extension GL_EXT_shader_8bit_storage : require -#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require -#extension GL_EXT_control_flow_attributes : enable - -#define PRECISION ${PRECISION} - -$if WEIGHT_NBITS == 4: - #define WEIGHT_INT4 - -$if HAS_BIAS: - #define HAS_BIAS - -$if WEIGHT_STORAGE == "buffer": - #define WEIGHT_BUFFER - -$if IO_STORAGE == "texture3d": - #define IO_TEXTURE - -layout(std430) buffer; - -#include "common.glslh" - -// Bindings — match add_linear_dqa_qw_node arg order: -// output(0), fp_input(1), packed_int8_input(2), int_input_sums(3 - unused), -// input_scales(4), input_zps(5), packed_weight(6), weight_sums(7), -// weight_scales(8), bias(9). -${layout_declare_tensor(B, "w", "t_output", "half", IO_STORAGE, is_scalar_array=True)} -// t_input is unread here -- the activations arrive already quantized in -// t_packed_int8_input -- but stays declared so the binding layout matches the -// dispatch site. It tracks IO_STORAGE so the two IO tensors stay consistent. -${layout_declare_tensor(B, "r", "t_input", "half", IO_STORAGE, is_scalar_array=False)} -// ROW-MAJOR (kPackedInt8_4W) packed activations, bound as a scalar int8_t -// array (row stride = K int8). Two things differ from dbuf4, which takes the -// 4h4w ivec4 block layout: -// 1. row-major, so a coopMatLoad can address it at all; -// 2. element type int8_t, MATCHING the coopmat component type. -// (2) is not cosmetic. Binding the same memory as int[] and loading a -// coopmat from it -- a type mismatch that demonstrably works for the -// Workgroup storage class, which is how the MMA loop reads Ash_int8 below -- -// silently produces wrong results from a StorageBuffer on this driver. -// The reference shmem_double_buf4-tr.comp sidesteps it the same way: its -// buffer_reference is declared `A_TYPE x[]`, i.e. int8_t for the int8 config. -// All A offsets/strides here are therefore in INT8 elements, not int. -${layout_declare_tensor(B, "r", "t_packed_int8_input", "int8", "buffer", is_scalar_array=True)} -${layout_declare_tensor(B, "r", "t_int8_input_sums", "int", "buffer", is_scalar_array=True)} -${layout_declare_tensor(B, "r", "t_int8_input_scales", "half", "texture3d")} -${layout_declare_tensor(B, "r", "t_int8_input_zps", "int8", "texture3d")} -${layout_declare_tensor(B, "r", "t_packed_weight", "int", WEIGHT_STORAGE, is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_weight_sums", "int", "buffer", is_scalar_array=True)} -${layout_declare_tensor(B, "r", "t_weight_scales", "half", "buffer", is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_bias", "half", "buffer", is_scalar_array=True)} - -${layout_declare_ubo(B, "ivec4", "output_sizes")} -${layout_declare_ubo(B, "ivec4", "input_sizes")} - -layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; - -${layout_declare_spec_const(C, "int", "apply_bias", "0")} -// INT4 only; inert (0) for INT8 so the dispatcher's spec list lines up. -${layout_declare_spec_const(C, "int", "K4_per_group", "0")} -${layout_declare_spec_const(C, "int", "num_groups_arg", "0")} -${layout_declare_spec_const(C, "int", "out_N_arg", "0")} - -// Tile geometry -const uint MMA_M = ${MMA_M}; -const uint MMA_N = ${MMA_N}; -const uint MMA_K = ${MMA_K}; - -const uint WG_TILE_M = ${WG_TILE_M}; -const uint WG_TILE_N = ${WG_TILE_N}; -const uint WG_TILE_K = ${WG_TILE_K}; - -const uint SG_GRID_X = ${SG_GRID_X}; -const uint SG_GRID_Y = ${SG_GRID_Y}; -const uint SUBGROUP_SIZE = ${SUBGROUP_SIZE}; -const uint NUM_SUBGROUPS = SG_GRID_X * SG_GRID_Y; -const uint WG_SIZE = NUM_SUBGROUPS * SUBGROUP_SIZE; - -const uint SG_TILE_M = WG_TILE_M / SG_GRID_Y; -const uint SG_TILE_N = WG_TILE_N / SG_GRID_X; -const uint MMAS_PER_SG_M = SG_TILE_M / MMA_M; -const uint MMAS_PER_SG_N = SG_TILE_N / MMA_N; - -const uint A_SLAB_INT8 = WG_TILE_M * MMA_K; -const uint B_USEFUL_U32 = MMA_K / 4u; -const uint B_STRIDE_U32 = B_USEFUL_U32 + 1u; // +1 skew -const uint B_SLAB_U32 = WG_TILE_N * B_STRIDE_U32; -const uint NUM_K_SLABS = WG_TILE_K / MMA_K; - -const uint A_SLAB_U32 = A_SLAB_INT8 / 4u; -const uint A_STRIDE_U32 = MMA_K / 4u; - -// One ping-pong slice covers all K-slabs of one chunk. -const uint ASH_SLICE_U32 = NUM_K_SLABS * A_SLAB_U32; -const uint BSH_SLICE_U32 = NUM_K_SLABS * B_SLAB_U32; - -// Double-buffered MMA operand staging. -shared uint Ash_int8[2u * ASH_SLICE_U32]; -shared uint Bsh_int8[2u * BSH_SLICE_U32]; - -// Per-WG-tile-row activation params (loaded ONCE at WG start; constant -// across groups). -shared int izp_sh[WG_TILE_M]; // int32 (cast from int8 source) for broadcast -shared float ifs_sh[WG_TILE_M]; // float32 (cast from fp16 source) for broadcast - -// Per-(group, output-channel) weight params, ping-ponged by group parity. -// (For per-channel INT8 only slice 0 is ever used.) -shared int wsum_sh[2u * WG_TILE_N]; -shared float wsc_sh[2u * WG_TILE_N]; - -#ifdef HAS_BIAS -shared float bias_sh[WG_TILE_N]; -#endif - -#ifdef IO_TEXTURE -// Result staging for the imageStore epilogue, mirroring the fp16 kernel: -// SG_GRID_Y bands of MMA_M rows, each WG_TILE_N wide, row-major. A full -// WG_TILE_M x WG_TILE_N buffer would cost SG_GRID_Y/MMAS_PER_SG_M x more LDS -// and wreck occupancy. float16_t-typed because coopMatStore needs it. -const uint CSH_ROWS = SG_GRID_Y * MMA_M; -shared float16_t Csh_out[CSH_ROWS * WG_TILE_N]; -#endif - -// Running fp32 accumulator (across all groups). -coopmat - result[MMAS_PER_SG_M][MMAS_PER_SG_N]; - -// Per-group int32 MMA accumulator. -coopmat - accum_int32[MMAS_PER_SG_M][MMAS_PER_SG_N]; - -void main() { - const uvec2 tileID = uvec2(gl_WorkGroupID.xy); - const uvec2 warpInTile = uvec2( - gl_SubgroupID % SG_GRID_X, - gl_SubgroupID / SG_GRID_X); - - const uint K = uint(input_sizes.x); - const uint N = uint(output_sizes.x); - const uint N4 = (N + 3u) / 4u; - const uint nblocks_x_A = (K + 3u) >> 2u; - -#ifdef WEIGHT_INT4 - const uint num_groups = uint(num_groups_arg); - const uint CHUNKS_PER_GROUP = uint(K4_per_group) * 4u / WG_TILE_K; -#else - // Per-channel: a single quant "group" spanning all of K. The nested - // groups x chunks loop below collapses to a flat chunk loop, the wsum/wsc - // ping-pong never crosses a boundary, and the epilog runs exactly once. - const uint num_groups = 1u; - const uint CHUNKS_PER_GROUP = uint(num_groups_arg); -#endif - const uint num_chunks = num_groups * CHUNKS_PER_GROUP; - - const uint tile_m_start = WG_TILE_M * tileID.y; - const uint tile_n_start = WG_TILE_N * tileID.x; - - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - result[i][j] = coopmat(0.0); - accum_int32[i][j] = coopmat(0); - } - } - - const uint K_BLOCKS_PER_CHUNK = WG_TILE_K >> 2u; - - // --- A staging tile map: one MMA_M x MMA_K coopmat tile per subgroup per - // slot. A chunk holds A_TILES_M x A_TILES_K such tiles; they are dealt - // round-robin across the NUM_SUBGROUPS subgroups, so every subgroup - // participates (dbuf4's per-thread map leaves WG_SIZE - - // A_ACTIVE_THREADS invocations idle whenever the tile is small). - // A_TILES_PER_SG rounds up, so the last slot may be partially used -- - // the `t < NUM_A_TILES` guard below is subgroup-uniform (t depends only - // on gl_SubgroupID), which is what coopmat ops require. - const uint A_TILES_M = WG_TILE_M / MMA_M; - const uint A_TILES_K = WG_TILE_K / MMA_K; // == NUM_K_SLABS - const uint NUM_A_TILES = A_TILES_M * A_TILES_K; - const uint A_TILES_PER_SG = (NUM_A_TILES + NUM_SUBGROUPS - 1u) / NUM_SUBGROUPS; - -#ifdef WEIGHT_INT4 - // --- B staging thread map: (block, col) slots; each slot extracts one - // ColumnMajor LDS uint (4 K-contiguous sign-extended int8) --- - const uint B_TOTAL_SLOTS = K_BLOCKS_PER_CHUNK * WG_TILE_N; - const uint B_SLOTS_PER_THREAD = B_TOTAL_SLOTS / WG_SIZE; - const uint N8_PER_TILE = WG_TILE_N >> 3u; -#else - // --- B staging thread map: one (k4, n4) ivec4 block per active thread --- - const uint B_FETCH_SLOTS = K_BLOCKS_PER_CHUNK * (WG_TILE_N >> 2u); - const uint N4_PER_TILE = WG_TILE_N >> 2u; - const uint b_k4_in_chunk = gl_LocalInvocationID.x / N4_PER_TILE; - const uint b_n_uint_col = gl_LocalInvocationID.x % N4_PER_TILE; - const bool b_active = gl_LocalInvocationID.x < B_FETCH_SLOTS; -#endif - - // Prefetch temp registers. temp_A is a coopmat array (the -tr change); - // indices into it are [[unroll]]-resolved compile-time constants, never - // dynamic -- dynamic indexing of a coopmat array is exactly the construct - // the Xclipse/AMD-PAL compiler has miscompiled before. - // (temp_A removed) -#ifdef WEIGHT_INT4 - ivec4 temp_B[B_SLOTS_PER_THREAD]; - int temp_wsum; - float temp_wsc; -#else - ivec4 temp_B; -#endif - - // ========================================================= - // PROLOGUE - // ========================================================= - if (gl_LocalInvocationID.x < (WG_TILE_M >> 2u)) { - const uint m4 = (tile_m_start >> 2u) + gl_LocalInvocationID.x; - const vec4 sc = vec4(texelFetch(t_int8_input_scales, ivec3(m4, 0, 0), 0)); - const ivec4 zp = texelFetch(t_int8_input_zps, ivec3(m4, 0, 0), 0); - const uint base = gl_LocalInvocationID.x * 4u; - ifs_sh[base + 0u] = sc.x; ifs_sh[base + 1u] = sc.y; - ifs_sh[base + 2u] = sc.z; ifs_sh[base + 3u] = sc.w; - izp_sh[base + 0u] = zp.x; izp_sh[base + 1u] = zp.y; - izp_sh[base + 2u] = zp.z; izp_sh[base + 3u] = zp.w; - } - // Group 0 weight sums/scales -> slice 0. - if (gl_LocalInvocationID.x < WG_TILE_N) { - const uint n_idx = tile_n_start + gl_LocalInvocationID.x; - f16vec4 sv = t_weight_scales[n_idx >> 2u]; - wsc_sh[gl_LocalInvocationID.x] = float(sv[n_idx & 3u]); - wsum_sh[gl_LocalInvocationID.x] = t_weight_sums[n_idx]; - } - memoryBarrierShared(); - barrier(); - - // izp/ifs are per-row activation params, constant across K groups — - // broadcast them into coopmats ONCE; the group epilog reuses them every - // group (they depend only on the row block i, not on the group or j). - coopmat - izp_bcast[MMAS_PER_SG_M]; - coopmat - ifs_bcast[MMAS_PER_SG_M]; - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - const uint local_m_base = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); - coopMatLoad( - izp_bcast[i], izp_sh, - local_m_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutColumnMajor); - coopMatLoad( - ifs_bcast[i], ifs_sh, - local_m_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutColumnMajor); - } - - // dbuf4: prefetch chunk 0 into temp registers, THEN store to slice 0 (no - // barrier here -- the main loop's first iteration barriers before - // reading slice 0). - // (A prefetch removed: loaded directly in the MMA loop) -#ifdef WEIGHT_INT4 - [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { - const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; - const uint block_in_chunk = slot >> 3u; - const uint k4_blk = block_in_chunk / N8_PER_TILE; - const uint n8_blk = (tile_n_start >> 3u) + (block_in_chunk % N8_PER_TILE); -#ifdef WEIGHT_BUFFER - temp_B[si] = t_packed_weight[(n8_blk * nblocks_x_A) + k4_blk]; -#else - temp_B[si] = texelFetch(t_packed_weight, ivec2(k4_blk, n8_blk), 0); -#endif - } -#else - if (b_active) { - const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; -#ifdef WEIGHT_BUFFER - temp_B = t_packed_weight[(b_k4_in_chunk * N4) + block_x_w]; -#else - temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, b_k4_in_chunk), 0); -#endif - } -#endif - { - // store chunk 0 -> slice 0 - // (A LDS store removed) -#ifdef WEIGHT_INT4 - [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { - const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; - const uint block_in_chunk = slot >> 3u; - const uint col_in_block = slot & 7u; - const uint k4_in_chunk = block_in_chunk / N8_PER_TILE; - const uint n8_in_tile = block_in_chunk % N8_PER_TILE; - const uint r = col_in_block & 3u; - const uint parity = col_in_block >> 2u; - const int w = temp_B[si][r]; - const int base = int(4u * parity); - const int v0 = (((w >> (base + 0)) & 0xF) - 8) & 0xFF; - const int v1 = (((w >> (base + 8)) & 0xF) - 8) & 0xFF; - const int v2 = (((w >> (base + 16)) & 0xF) - 8) & 0xFF; - const int v3 = (((w >> (base + 24)) & 0xF) - 8) & 0xFF; - const uint n_col = n8_in_tile * 8u + r + parity * 4u; - const uint slab_idx = k4_in_chunk / (MMA_K >> 2u); - const uint k4_in_slab = k4_in_chunk % (MMA_K >> 2u); - Bsh_int8[slab_idx * B_SLAB_U32 + n_col * B_STRIDE_U32 + k4_in_slab] = - uint(v0 | (v1 << 8) | (v2 << 16) | (v3 << 24)); - } -#else - if (b_active) { - const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); - const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); - const uint n_col_base = b_n_uint_col * 4u; - [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { - Bsh_int8[slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab] = - uint(temp_B[n_in_blk]); - } - } -#endif - } - - // ========================================================= - // MAIN LOOP (dbuf4) — nested groups x chunks (kept nested; flattening it - // with a conditional coopmat epilog crashes the Xclipse PAL compiler at - // large spec-resolved trip counts). One barrier per chunk. Chunk - // iteration (global index `chunk`): - // 1. barrier — A/B slice (chunk%2) fully written; on the first chunk - // of group g, wsum/wsc slice (g%2) is too. - // 2. prefetch — chunk+1 (A blocks, B blocks) into temp; when chunk+1 - // starts a new group, also its wsum/wsc element. Skipped - // entirely on the final chunk. - // 3. int8 MMA — on slice (chunk%2) into accum_int32. - // 4. store — temp -> A/B slice ((chunk+1)%2), unpacking the weight; - // on a group boundary, wsum/wsc -> slice ((g+1)%2). - // The group epilog runs unconditionally at the tail of each group. - // ========================================================= - uint chunk = 0; - for (uint group_i = 0; group_i < num_groups; ++group_i) { - for (uint inner = 0; inner < CHUNKS_PER_GROUP; ++inner, ++chunk) { - const bool has_next = chunk + 1u < num_chunks; - const bool group_crossing = has_next && (inner + 1u == CHUNKS_PER_GROUP); - const uint cur_a = (chunk % 2u) * ASH_SLICE_U32; - const uint cur_b = (chunk % 2u) * BSH_SLICE_U32; - const uint nxt_a = ((chunk + 1u) % 2u) * ASH_SLICE_U32; - const uint nxt_b = ((chunk + 1u) % 2u) * BSH_SLICE_U32; - - barrier(); - - // --- 2. prefetch chunk+1 -> temp --- - if (has_next) { - const uint chunkK_nxt = (chunk + 1u) * WG_TILE_K; - // (A prefetch removed) -#ifdef WEIGHT_INT4 - [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { - const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; - const uint block_in_chunk = slot >> 3u; - const uint k4_blk = (chunkK_nxt >> 2u) + block_in_chunk / N8_PER_TILE; - const uint n8_blk = (tile_n_start >> 3u) + (block_in_chunk % N8_PER_TILE); -#ifdef WEIGHT_BUFFER - temp_B[si] = t_packed_weight[(n8_blk * nblocks_x_A) + k4_blk]; -#else - temp_B[si] = texelFetch(t_packed_weight, ivec2(k4_blk, n8_blk), 0); -#endif - } - if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { - const uint n_idx = tile_n_start + gl_LocalInvocationID.x; - f16vec4 sv = t_weight_scales[(group_i + 1u) * N4 + (n_idx >> 2u)]; - temp_wsc = float(sv[n_idx & 3u]); - temp_wsum = t_weight_sums[(group_i + 1u) * N + n_idx]; - } -#else - if (b_active) { - const uint block_y_w = (chunkK_nxt >> 2u) + b_k4_in_chunk; - const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; -#ifdef WEIGHT_BUFFER - temp_B = t_packed_weight[(block_y_w * N4) + block_x_w]; -#else - temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, block_y_w), 0); -#endif - } -#endif - } - - // --- 3. int8 MMA on the cur slice --- - [[unroll]] for (uint k = 0; k < NUM_K_SLABS; ++k) { - const uint slab_a_base_u32 = cur_a + k * A_SLAB_U32; - const uint slab_b_base_u32 = cur_b + k * B_SLAB_U32; - - coopmat matA[MMAS_PER_SG_M]; - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - const uint row_a = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); - // Offsets/strides are in int8 elements, matching the int8_t - // binding. With A bound as int[] instead, this load silently - // produced wrong results regardless of which unit was used. - const uint a_row_stride_i8 = nblocks_x_A * 4u; // int8 elements - coopMatLoad( - matA[i], t_packed_int8_input, - (tile_m_start + row_a) * a_row_stride_i8 + - (chunk * WG_TILE_K + k * MMA_K), - a_row_stride_i8, - gl_CooperativeMatrixLayoutRowMajor); - } - - coopmat matB; - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint col_b = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopMatLoad( - matB, Bsh_int8, - slab_b_base_u32 + col_b * B_STRIDE_U32, - B_STRIDE_U32, - gl_CooperativeMatrixLayoutColumnMajor); - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - accum_int32[i][j] = coopMatMulAdd(matA[i], matB, accum_int32[i][j]); - } - } - } - - // --- 4. store temp (chunk+1) -> nxt slice --- - if (has_next) { - // (A LDS store removed) -#ifdef WEIGHT_INT4 - [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { - const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; - const uint block_in_chunk = slot >> 3u; - const uint col_in_block = slot & 7u; - const uint k4_in_chunk = block_in_chunk / N8_PER_TILE; - const uint n8_in_tile = block_in_chunk % N8_PER_TILE; - const uint r = col_in_block & 3u; - const uint parity = col_in_block >> 2u; - const int w = temp_B[si][r]; - const int base = int(4u * parity); - const int v0 = (((w >> (base + 0)) & 0xF) - 8) & 0xFF; - const int v1 = (((w >> (base + 8)) & 0xF) - 8) & 0xFF; - const int v2 = (((w >> (base + 16)) & 0xF) - 8) & 0xFF; - const int v3 = (((w >> (base + 24)) & 0xF) - 8) & 0xFF; - const uint n_col = n8_in_tile * 8u + r + parity * 4u; - const uint slab_idx = k4_in_chunk / (MMA_K >> 2u); - const uint k4_in_slab = k4_in_chunk % (MMA_K >> 2u); - Bsh_int8[nxt_b + slab_idx * B_SLAB_U32 + n_col * B_STRIDE_U32 + k4_in_slab] = - uint(v0 | (v1 << 8) | (v2 << 16) | (v3 << 24)); - } - if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { - const uint wbase_nxt = ((group_i + 1u) % 2u) * WG_TILE_N; - wsum_sh[wbase_nxt + gl_LocalInvocationID.x] = temp_wsum; - wsc_sh[wbase_nxt + gl_LocalInvocationID.x] = temp_wsc; - } -#else - if (b_active) { - const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); - const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); - const uint n_col_base = b_n_uint_col * 4u; - [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { - Bsh_int8[nxt_b + slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab] = - uint(temp_B[n_in_blk]); - } - } -#endif - } - } // chunks - - // --- Group epilog: dequant accum_int32 -> result, reset accum --- - { - const uint wbase = (group_i % 2u) * WG_TILE_N; - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - - coopmat wsum_bcast; - coopMatLoad( - wsum_bcast, wsum_sh, - wbase + local_n_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutRowMajor); - - coopmat wsc_bcast; - coopMatLoad( - wsc_bcast, wsc_sh, - wbase + local_n_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutRowMajor); - - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - coopmat adjusted = - accum_int32[i][j] - izp_bcast[i] * wsum_bcast; - coopmat adjusted_fp = - coopmat(adjusted); - coopmat scales_outer = - ifs_bcast[i] * wsc_bcast; - result[i][j] += adjusted_fp * scales_outer; - accum_int32[i][j] = coopmat(0); - } - } - } - } // groups - - // --- Bias (optional) --- -#ifdef HAS_BIAS - if (apply_bias > 0) { - for (uint t = gl_LocalInvocationID.x; t < WG_TILE_N; t += WG_SIZE) { - bias_sh[t] = float(t_bias[tile_n_start + t]); - } - memoryBarrierShared(); - barrier(); - } -#endif - - // --- Store result tile --- - // N for the store address math MUST come from the spec constant, not the - // sizes UBO (see out_N_arg above). -#ifdef IO_TEXTURE - // Epilogue iteration i drains accumulator row-block i from EVERY subgroup - // into Csh_out at once, so the SG_GRID_Y bands it holds are disjoint global - // row ranges; the whole workgroup then imageStores them. lr / MMA_M is the - // writing subgroup's warpInTile.y, so the global row reproduces the buffer - // path's gi exactly. - // - // PORTABILITY NOTE: the barrier() in the loop body keeps this loop rolled - // despite [[unroll]], so result[i][j] IS dynamically indexed. Coopmat arrays - // are opaque per-lane storage and dynamic indexing is exactly the construct - // the Xclipse/AMD-PAL compiler has broken before -- check this first if the - // texture variants miscompile on M51. - const uint CSH_TEXELS_PER_ROW = WG_TILE_N / 4u; - const uint CSH_TEXELS = CSH_ROWS * CSH_TEXELS_PER_ROW; - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - // Guards Csh_out against the previous iteration's readers. Inert on i == 0 - // but must stay unconditional to remain workgroup-uniform. - barrier(); - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { -#ifdef HAS_BIAS - if (apply_bias > 0) { - const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopmat bias_tile; - coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); - result[i][j] += bias_tile; - } -#endif - coopmat out_tile = - coopmat(result[i][j]); - coopMatStore( - out_tile, Csh_out, - warpInTile.y * MMA_M * WG_TILE_N + - MMA_N * (MMAS_PER_SG_N * warpInTile.x + j), - WG_TILE_N, - gl_CooperativeMatrixLayoutRowMajor); - } - memoryBarrierShared(); - barrier(); - - for (uint t = gl_LocalInvocationID.x; t < CSH_TEXELS; t += WG_SIZE) { - const uint lr = t / CSH_TEXELS_PER_ROW; - const uint lc4 = t % CSH_TEXELS_PER_ROW; - const uint m = - tile_m_start + (lr / MMA_M) * SG_TILE_M + i * MMA_M + (lr % MMA_M); - const uint base = lr * WG_TILE_N + lc4 * 4u; - imageStore( - t_output, - ivec3(tile_n_start / 4u + lc4, m, 0), - vec4( - float(Csh_out[base]), - float(Csh_out[base + 1u]), - float(Csh_out[base + 2u]), - float(Csh_out[base + 3u]))); - } - } -#else - const uint N_out = uint(out_N_arg); - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint gi = tile_m_start + MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); - const uint gj = tile_n_start + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - -#ifdef HAS_BIAS - if (apply_bias > 0) { - const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopmat bias_tile; - coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); - result[i][j] += bias_tile; - } -#endif - - coopmat out_tile = - coopmat(result[i][j]); - coopMatStore( - out_tile, t_output, - gi * N_out + gj, N_out, - gl_CooperativeMatrixLayoutRowMajor); - } - } -#endif // IO_TEXTURE -} diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trd.yaml b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trd.yaml deleted file mode 100644 index e26c6c3bea6..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trd.yaml +++ /dev/null @@ -1,188 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -# "-tr" (coopmat-staged A) variants of the int8 dq8ca_q4gsw coopmat kernel, -# ported from shmem_double_buf4-tr.comp (vk_cooperative_matrix_perf, -# gemm-ubm). linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trd.glsl is a fork of -# linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4.glsl with ONLY the A-side global -> -# LDS staging swapped from a per-thread ivec4 copy to -# coopMatLoad(global) -> coopmat<> -> coopMatStore(LDS). B staging is -# byte-identical (int4 nibble unpack cannot be coopmat-staged) and B was -# already column-major in LDS, so it is not a delta. -# -# These variants require t_packed_int8_input in the ROW-MAJOR -# kPackedInt8_4W layout (produced by quantize_and_pack_4w_with_group_sums), -# NOT the stock 4h4w block layout. QuantizedLinear.cpp switches the tensor -# layout and the packer node together with the variant token, and only when -# the coopmat gate passes -- see dq8ca_wants_rowmajor_int8_input(). -# -# Selected via -# ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4trd_txkgs. -# NOT the shipped default: unvalidated until it passes repeated -# test_llama_microbench --correctness-only runs (a single pass is not proof -- -# see dq8ca_coopmat_variant()'s comment). -# -# Tile-geometry preconditions beyond dbuf4's: WG_TILE_M % MMA_M == 0 and -# WG_TILE_K % MMA_K == 0 (both hold for every entry below). The seed set is -# deliberately small -- widen it from a sweep once correctness is established. - -linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trd: - parameter_names_with_default_values: - PRECISION: highp - HAS_BIAS: false - IO_STORAGE: buffer - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - MMA_M: 16 - MMA_N: 16 - MMA_K: 16 - WG_TILE_M: 64 - WG_TILE_N: 32 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - shader_variants: - # production 8da4w tile (dbuf2/dbuf4 default) -- the A/B anchor - # A tiles/chunk = 8, subgroups = 2 -> A_TILES_PER_SG = 4 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trd_t64x32k32g12s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 32 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trd_t64x32k32g12s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 32 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trd_t64x32k32g12s64_buffer_buffer_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: buffer - WG_TILE_M: 64 - WG_TILE_N: 32 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - # q4gsw production geometry; 8 A-tiles over 4 subgroups - # A tiles/chunk = 8, subgroups = 4 -> A_TILES_PER_SG = 2 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trd_t128x128k16g22s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 2 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trd_t128x128k16g22s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 2 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - # 1 A-tile per subgroup (A_TILES_PER_SG == 1, no leftover slot) - # A tiles/chunk = 8, subgroups = 8 -> A_TILES_PER_SG = 1 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trd_t64x64k32g24s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trd_t64x64k32g24s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - # 16 A-tiles over 8 subgroups - # A tiles/chunk = 16, subgroups = 8 -> A_TILES_PER_SG = 2 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trd_t128x64k32g24s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trd_t128x64k32g24s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - # Reference-matching tile: shmem_double_buf4-tr.comp's own geometry - # (TILE 128x128x16, WORKGROUP_WIDTH_IN_SUBGROUPS=4 x DBUF4_GRID_HEIGHT=2, - # wave32) => 8 subgroups / 256 threads, MMAS_PER_SG 4x2. Added so our - # kernel can be A/B'd against the teammate's PAL capture at the same - # tile and the same 2048x1024x4096 shape (8b wk_wv prefill). - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trd_t128x128k16g42s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trd_t128x128k16g42s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - # G2/Pavan's documented INT8 uBM config: TILE 128x128x64 with a 4x2 - # subgroup grid at wave32 (workgroupSize=256). Added to compare our kernel - # against their reference dbuf4 number (~1973us) at the same geometry and - # the same shape (M2048 K4096 N1024 == 8b wk_wv prefill). - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trd_t128x128k64g42s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trd_t128x128k64g42s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trm.glsl b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trm.glsl deleted file mode 100644 index b12d6c21012..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trm.glsl +++ /dev/null @@ -1,707 +0,0 @@ -/* - * Copyright (c) Meta Platforms, Inc. and affiliates. - * All rights reserved. - * - * This source code is licensed under the BSD-style license found in the - * LICENSE file in the root directory of this source tree. - */ - -/* - * DIAGNOSTIC BISECT variant of linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr. - * - * Identical to dbuf4tr -- same ROW-MAJOR kPackedInt8_4W activation layout, same - * quantize_and_pack_4w_with_group_sums producer, same LDS layout, same loop -- - * EXCEPT that A staging is done with plain scalar loads/stores instead of - * coopMatLoad/coopMatStore. It exists to split a dbuf4tr correctness failure - * into "the row-major packer/layout is wrong" (this variant also fails) vs - * "the coopmat staging is wrong" (this variant passes). - * - * Not a perf candidate; delete once dbuf4tr is understood. - * - * Original dbuf4tr header follows. - * - * "-tr" (coopmat-staged A) variant of linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4. - * - * Ported from shmem_double_buf4-tr.comp on vk_cooperative_matrix_perf's - * gemm-ubm branch. That reference file's delta over shmem_double_buf4.comp is - * that the global -> LDS staging goes through COOPERATIVE MATRIX REGISTERS - * (coopMatLoad from global -> coopmat<> array -> coopMatStore into shared) - * instead of a hand-rolled per-thread uvec4 copy, and that B lands in LDS - * column-major. - * - * Only the A half of that idea is portable to this kernel: - * - * - B is ALREADY column-major in LDS here (Bsh_int8 is K-contiguous per - * output column, read back with gl_CooperativeMatrixLayoutColumnMajor), - * so the reference's "tr" property is not a delta for B at all. - * - B CANNOT be coopmat-staged: the weights are int4, each ivec4 holding - * 8 columns x 4 K-values that need the nibble-extract / -8 / sign-pack - * below. coopMatLoad cannot unpack nibbles, and a coopmat's per-lane - * layout is opaque so one cannot be built from unpacked registers. - * B staging is therefore left byte-identical to dbuf4. - * - A CAN be coopmat-staged, but only against a ROW-MAJOR packed int8 - * activation buffer. The stock 4h4w layout (kPackedInt8_4H4W, produced by - * quantize_and_pack_4h4w_with_group_sums.glsl) is NOT row-major: element - * [m4 * K4 + k4] is an ivec4 whose COMPONENT selects one of 4 rows, so as - * a uint array the index is m4*(4*K4) + k4*4 + r, which is not affine in - * the row index and cannot be addressed by any RowMajor/ColumnMajor - * coopMatLoad. (ColumnMajor is out on contiguity too: a uint packs 4 - * K-values, not 4 M-values.) - * - * So this shader binds t_packed_int8_input as a SCALAR int array in the - * kPackedInt8_4W layout -- plain row-major int8, 4 K-values per int32, - * row stride K4 -- produced by quantize_and_pack_4w_with_group_sums.glsl. - * QuantizedLinear.cpp allocates that layout (and dispatches that packer) - * only when the active dq8ca variant is a "tsweep_dbuf4tr_t..." token AND - * the coopmat gate passes, so the tiled fallback never sees the wrong - * layout. Everything downstream of A staging -- LDS layout, int8 WMMA thread - * maps, group epilog, bias/store epilogue -- is unchanged from dbuf4. - * - * A staging (the actual -tr port): - * dbuf4: per-thread (m4, k4) ivec4 fetch; only A_ACTIVE_THREADS = - * (WG_TILE_M/4) * (WG_TILE_K/4) invocations participate, each - * scattering 4 rows into Ash_int8 with 4 scalar stores. - * dbuf4tr: per-SUBGROUP MMA_M x MMA_K tile fetch via coopMatLoad straight - * from global, then coopMatStore into the same Ash_int8 slot. - * The (WG_TILE_M/MMA_M) * (WG_TILE_K/MMA_K) tiles of a chunk are - * dealt round-robin across the NUM_SUBGROUPS subgroups. - * - * The loop structure is dbuf4's, unchanged: - * prologue: prefetch chunk 0 -> temp, store to slice 0 (no barrier) - * per iter: barrier -> prefetch(next) -> MMA(cur) -> store(next) - * and the nested `groups x chunks` loop with an unconditional group epilog is - * kept as-is (flattening it crashes the Xclipse PAL compiler at large - * spec-resolved trip counts -- see dbuf2's header). - * - * Selected at dispatch via - * ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4tr_txkgs<32|64> - * (QuantizedLinear.cpp), additive to the tsweep_dbuf4_t... and tsweep_t... - * namespaces. NOT the default -- unvalidated until it passes repeated - * test_llama_microbench --correctness-only runs (see dq8ca_coopmat_variant()'s - * comment on why a single pass is not proof). - * - * Performs: out[M,N] = dequant(int8_act) * dequant(int_w) (+ bias) - * via coopmat x coopmat -> coopmat on the matrix unit. - * - * Hard preconditions (in addition to dbuf4's): - * M % WG_TILE_M == 0, N % WG_TILE_N == 0, K % WG_TILE_K == 0, - * group_size % WG_TILE_K == 0, K % 4 == 0, - * WG_TILE_M % MMA_M == 0, WG_TILE_K % MMA_K == 0, - * t_packed_int8_input in kPackedInt8_4W (row-major) layout, - * device exposes coopmatx-> at 16x16x16. - */ - -#version 450 core - -#extension GL_KHR_cooperative_matrix : require -#extension GL_KHR_memory_scope_semantics : require -#extension GL_KHR_shader_subgroup_basic : enable -#extension GL_EXT_shader_explicit_arithmetic_types : require -#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require -#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require -#extension GL_EXT_control_flow_attributes : enable - -#define PRECISION ${PRECISION} - -$if WEIGHT_NBITS == 4: - #define WEIGHT_INT4 - -$if HAS_BIAS: - #define HAS_BIAS - -$if WEIGHT_STORAGE == "buffer": - #define WEIGHT_BUFFER - -$if IO_STORAGE == "texture3d": - #define IO_TEXTURE - -layout(std430) buffer; - -#include "common.glslh" - -// Bindings — match add_linear_dqa_qw_node arg order: -// output(0), fp_input(1), packed_int8_input(2), int_input_sums(3 - unused), -// input_scales(4), input_zps(5), packed_weight(6), weight_sums(7), -// weight_scales(8), bias(9). -${layout_declare_tensor(B, "w", "t_output", "half", IO_STORAGE, is_scalar_array=True)} -// t_input is unread here -- the activations arrive already quantized in -// t_packed_int8_input -- but stays declared so the binding layout matches the -// dispatch site. It tracks IO_STORAGE so the two IO tensors stay consistent. -${layout_declare_tensor(B, "r", "t_input", "half", IO_STORAGE, is_scalar_array=False)} -// ROW-MAJOR (kPackedInt8_4W) packed activations: scalar int array, each -// element holding 4 K-contiguous int8, row stride K4 = K/4. This is the one -// binding that differs from dbuf4 (which takes the 4h4w ivec4 block layout); -// it is what makes the coopMatLoad-based A staging below addressable. -${layout_declare_tensor(B, "r", "t_packed_int8_input", "int", "buffer", is_scalar_array=True)} -${layout_declare_tensor(B, "r", "t_int8_input_sums", "int", "buffer", is_scalar_array=True)} -${layout_declare_tensor(B, "r", "t_int8_input_scales", "half", "texture3d")} -${layout_declare_tensor(B, "r", "t_int8_input_zps", "int8", "texture3d")} -${layout_declare_tensor(B, "r", "t_packed_weight", "int", WEIGHT_STORAGE, is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_weight_sums", "int", "buffer", is_scalar_array=True)} -${layout_declare_tensor(B, "r", "t_weight_scales", "half", "buffer", is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_bias", "half", "buffer", is_scalar_array=True)} - -${layout_declare_ubo(B, "ivec4", "output_sizes")} -${layout_declare_ubo(B, "ivec4", "input_sizes")} - -layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; - -${layout_declare_spec_const(C, "int", "apply_bias", "0")} -// INT4 only; inert (0) for INT8 so the dispatcher's spec list lines up. -${layout_declare_spec_const(C, "int", "K4_per_group", "0")} -${layout_declare_spec_const(C, "int", "num_groups_arg", "0")} -${layout_declare_spec_const(C, "int", "out_N_arg", "0")} - -// Tile geometry -const uint MMA_M = ${MMA_M}; -const uint MMA_N = ${MMA_N}; -const uint MMA_K = ${MMA_K}; - -const uint WG_TILE_M = ${WG_TILE_M}; -const uint WG_TILE_N = ${WG_TILE_N}; -const uint WG_TILE_K = ${WG_TILE_K}; - -const uint SG_GRID_X = ${SG_GRID_X}; -const uint SG_GRID_Y = ${SG_GRID_Y}; -const uint SUBGROUP_SIZE = ${SUBGROUP_SIZE}; -const uint NUM_SUBGROUPS = SG_GRID_X * SG_GRID_Y; -const uint WG_SIZE = NUM_SUBGROUPS * SUBGROUP_SIZE; - -const uint SG_TILE_M = WG_TILE_M / SG_GRID_Y; -const uint SG_TILE_N = WG_TILE_N / SG_GRID_X; -const uint MMAS_PER_SG_M = SG_TILE_M / MMA_M; -const uint MMAS_PER_SG_N = SG_TILE_N / MMA_N; - -const uint A_SLAB_INT8 = WG_TILE_M * MMA_K; -const uint B_USEFUL_U32 = MMA_K / 4u; -const uint B_STRIDE_U32 = B_USEFUL_U32 + 1u; // +1 skew -const uint B_SLAB_U32 = WG_TILE_N * B_STRIDE_U32; -const uint NUM_K_SLABS = WG_TILE_K / MMA_K; - -const uint A_SLAB_U32 = A_SLAB_INT8 / 4u; -const uint A_STRIDE_U32 = MMA_K / 4u; - -// One ping-pong slice covers all K-slabs of one chunk. -const uint ASH_SLICE_U32 = NUM_K_SLABS * A_SLAB_U32; -const uint BSH_SLICE_U32 = NUM_K_SLABS * B_SLAB_U32; - -// Double-buffered MMA operand staging. -shared uint Ash_int8[2u * ASH_SLICE_U32]; -shared uint Bsh_int8[2u * BSH_SLICE_U32]; - -// Per-WG-tile-row activation params (loaded ONCE at WG start; constant -// across groups). -shared int izp_sh[WG_TILE_M]; // int32 (cast from int8 source) for broadcast -shared float ifs_sh[WG_TILE_M]; // float32 (cast from fp16 source) for broadcast - -// Per-(group, output-channel) weight params, ping-ponged by group parity. -// (For per-channel INT8 only slice 0 is ever used.) -shared int wsum_sh[2u * WG_TILE_N]; -shared float wsc_sh[2u * WG_TILE_N]; - -#ifdef HAS_BIAS -shared float bias_sh[WG_TILE_N]; -#endif - -#ifdef IO_TEXTURE -// Result staging for the imageStore epilogue, mirroring the fp16 kernel: -// SG_GRID_Y bands of MMA_M rows, each WG_TILE_N wide, row-major. A full -// WG_TILE_M x WG_TILE_N buffer would cost SG_GRID_Y/MMAS_PER_SG_M x more LDS -// and wreck occupancy. float16_t-typed because coopMatStore needs it. -const uint CSH_ROWS = SG_GRID_Y * MMA_M; -shared float16_t Csh_out[CSH_ROWS * WG_TILE_N]; -#endif - -// Running fp32 accumulator (across all groups). -coopmat - result[MMAS_PER_SG_M][MMAS_PER_SG_N]; - -// Per-group int32 MMA accumulator. -coopmat - accum_int32[MMAS_PER_SG_M][MMAS_PER_SG_N]; - -void main() { - const uvec2 tileID = uvec2(gl_WorkGroupID.xy); - const uvec2 warpInTile = uvec2( - gl_SubgroupID % SG_GRID_X, - gl_SubgroupID / SG_GRID_X); - - const uint K = uint(input_sizes.x); - const uint N = uint(output_sizes.x); - const uint N4 = (N + 3u) / 4u; - const uint nblocks_x_A = (K + 3u) >> 2u; - -#ifdef WEIGHT_INT4 - const uint num_groups = uint(num_groups_arg); - const uint CHUNKS_PER_GROUP = uint(K4_per_group) * 4u / WG_TILE_K; -#else - // Per-channel: a single quant "group" spanning all of K. The nested - // groups x chunks loop below collapses to a flat chunk loop, the wsum/wsc - // ping-pong never crosses a boundary, and the epilog runs exactly once. - const uint num_groups = 1u; - const uint CHUNKS_PER_GROUP = uint(num_groups_arg); -#endif - const uint num_chunks = num_groups * CHUNKS_PER_GROUP; - - const uint tile_m_start = WG_TILE_M * tileID.y; - const uint tile_n_start = WG_TILE_N * tileID.x; - - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - result[i][j] = coopmat(0.0); - accum_int32[i][j] = coopmat(0); - } - } - - const uint K_BLOCKS_PER_CHUNK = WG_TILE_K >> 2u; - - // --- A staging slot map (DIAGNOSTIC): one row-major int (4 K-contiguous - // int8 of ONE row) per thread per slot. Deliberately NOT coopmat. - const uint A_TOTAL_SLOTS = WG_TILE_M * K_BLOCKS_PER_CHUNK; - const uint A_SLOTS_PER_THREAD = (A_TOTAL_SLOTS + WG_SIZE - 1u) / WG_SIZE; - - // --- (unused in this variant) A staging tile map: one MMA_M x MMA_K coopmat tile per subgroup per - // slot. A chunk holds A_TILES_M x A_TILES_K such tiles; they are dealt - // round-robin across the NUM_SUBGROUPS subgroups, so every subgroup - // participates (dbuf4's per-thread map leaves WG_SIZE - - // A_ACTIVE_THREADS invocations idle whenever the tile is small). - // A_TILES_PER_SG rounds up, so the last slot may be partially used -- - // the `t < NUM_A_TILES` guard below is subgroup-uniform (t depends only - // on gl_SubgroupID), which is what coopmat ops require. - const uint A_TILES_M = WG_TILE_M / MMA_M; - const uint A_TILES_K = WG_TILE_K / MMA_K; // == NUM_K_SLABS - const uint NUM_A_TILES = A_TILES_M * A_TILES_K; - const uint A_TILES_PER_SG = (NUM_A_TILES + NUM_SUBGROUPS - 1u) / NUM_SUBGROUPS; - -#ifdef WEIGHT_INT4 - // --- B staging thread map: (block, col) slots; each slot extracts one - // ColumnMajor LDS uint (4 K-contiguous sign-extended int8) --- - const uint B_TOTAL_SLOTS = K_BLOCKS_PER_CHUNK * WG_TILE_N; - const uint B_SLOTS_PER_THREAD = B_TOTAL_SLOTS / WG_SIZE; - const uint N8_PER_TILE = WG_TILE_N >> 3u; -#else - // --- B staging thread map: one (k4, n4) ivec4 block per active thread --- - const uint B_FETCH_SLOTS = K_BLOCKS_PER_CHUNK * (WG_TILE_N >> 2u); - const uint N4_PER_TILE = WG_TILE_N >> 2u; - const uint b_k4_in_chunk = gl_LocalInvocationID.x / N4_PER_TILE; - const uint b_n_uint_col = gl_LocalInvocationID.x % N4_PER_TILE; - const bool b_active = gl_LocalInvocationID.x < B_FETCH_SLOTS; -#endif - - // Prefetch temp registers. temp_A is a coopmat array (the -tr change); - // indices into it are [[unroll]]-resolved compile-time constants, never - // dynamic -- dynamic indexing of a coopmat array is exactly the construct - // the Xclipse/AMD-PAL compiler has miscompiled before. - int temp_A[A_SLOTS_PER_THREAD]; -#ifdef WEIGHT_INT4 - ivec4 temp_B[B_SLOTS_PER_THREAD]; - int temp_wsum; - float temp_wsc; -#else - ivec4 temp_B; -#endif - - // ========================================================= - // PROLOGUE - // ========================================================= - if (gl_LocalInvocationID.x < (WG_TILE_M >> 2u)) { - const uint m4 = (tile_m_start >> 2u) + gl_LocalInvocationID.x; - const vec4 sc = vec4(texelFetch(t_int8_input_scales, ivec3(m4, 0, 0), 0)); - const ivec4 zp = texelFetch(t_int8_input_zps, ivec3(m4, 0, 0), 0); - const uint base = gl_LocalInvocationID.x * 4u; - ifs_sh[base + 0u] = sc.x; ifs_sh[base + 1u] = sc.y; - ifs_sh[base + 2u] = sc.z; ifs_sh[base + 3u] = sc.w; - izp_sh[base + 0u] = zp.x; izp_sh[base + 1u] = zp.y; - izp_sh[base + 2u] = zp.z; izp_sh[base + 3u] = zp.w; - } - // Group 0 weight sums/scales -> slice 0. - if (gl_LocalInvocationID.x < WG_TILE_N) { - const uint n_idx = tile_n_start + gl_LocalInvocationID.x; - f16vec4 sv = t_weight_scales[n_idx >> 2u]; - wsc_sh[gl_LocalInvocationID.x] = float(sv[n_idx & 3u]); - wsum_sh[gl_LocalInvocationID.x] = t_weight_sums[n_idx]; - } - memoryBarrierShared(); - barrier(); - - // izp/ifs are per-row activation params, constant across K groups — - // broadcast them into coopmats ONCE; the group epilog reuses them every - // group (they depend only on the row block i, not on the group or j). - coopmat - izp_bcast[MMAS_PER_SG_M]; - coopmat - ifs_bcast[MMAS_PER_SG_M]; - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - const uint local_m_base = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); - coopMatLoad( - izp_bcast[i], izp_sh, - local_m_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutColumnMajor); - coopMatLoad( - ifs_bcast[i], ifs_sh, - local_m_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutColumnMajor); - } - - // dbuf4: prefetch chunk 0 into temp registers, THEN store to slice 0 (no - // barrier here -- the main loop's first iteration barriers before - // reading slice 0). - [[unroll]] for (uint si = 0; si < A_SLOTS_PER_THREAD; ++si) { - const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; - if (slot < A_TOTAL_SLOTS) { - const uint row = slot / K_BLOCKS_PER_CHUNK; - const uint k4 = slot % K_BLOCKS_PER_CHUNK; - temp_A[si] = - t_packed_int8_input[(tile_m_start + row) * nblocks_x_A + k4]; - } - } -#ifdef WEIGHT_INT4 - [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { - const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; - const uint block_in_chunk = slot >> 3u; - const uint k4_blk = block_in_chunk / N8_PER_TILE; - const uint n8_blk = (tile_n_start >> 3u) + (block_in_chunk % N8_PER_TILE); -#ifdef WEIGHT_BUFFER - temp_B[si] = t_packed_weight[(n8_blk * nblocks_x_A) + k4_blk]; -#else - temp_B[si] = texelFetch(t_packed_weight, ivec2(k4_blk, n8_blk), 0); -#endif - } -#else - if (b_active) { - const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; -#ifdef WEIGHT_BUFFER - temp_B = t_packed_weight[(b_k4_in_chunk * N4) + block_x_w]; -#else - temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, b_k4_in_chunk), 0); -#endif - } -#endif - { - // store chunk 0 -> slice 0 - [[unroll]] for (uint si = 0; si < A_SLOTS_PER_THREAD; ++si) { - const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; - if (slot < A_TOTAL_SLOTS) { - const uint row = slot / K_BLOCKS_PER_CHUNK; - const uint k4 = slot % K_BLOCKS_PER_CHUNK; - const uint slab_idx = k4 / (MMA_K >> 2u); - const uint k_uint_in_slab = k4 % (MMA_K >> 2u); - Ash_int8 - [slab_idx * A_SLAB_U32 + row * A_STRIDE_U32 + k_uint_in_slab] = - uint(temp_A[si]); - } - } -#ifdef WEIGHT_INT4 - [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { - const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; - const uint block_in_chunk = slot >> 3u; - const uint col_in_block = slot & 7u; - const uint k4_in_chunk = block_in_chunk / N8_PER_TILE; - const uint n8_in_tile = block_in_chunk % N8_PER_TILE; - const uint r = col_in_block & 3u; - const uint parity = col_in_block >> 2u; - const int w = temp_B[si][r]; - const int base = int(4u * parity); - const int v0 = (((w >> (base + 0)) & 0xF) - 8) & 0xFF; - const int v1 = (((w >> (base + 8)) & 0xF) - 8) & 0xFF; - const int v2 = (((w >> (base + 16)) & 0xF) - 8) & 0xFF; - const int v3 = (((w >> (base + 24)) & 0xF) - 8) & 0xFF; - const uint n_col = n8_in_tile * 8u + r + parity * 4u; - const uint slab_idx = k4_in_chunk / (MMA_K >> 2u); - const uint k4_in_slab = k4_in_chunk % (MMA_K >> 2u); - Bsh_int8[slab_idx * B_SLAB_U32 + n_col * B_STRIDE_U32 + k4_in_slab] = - uint(v0 | (v1 << 8) | (v2 << 16) | (v3 << 24)); - } -#else - if (b_active) { - const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); - const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); - const uint n_col_base = b_n_uint_col * 4u; - [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { - Bsh_int8[slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab] = - uint(temp_B[n_in_blk]); - } - } -#endif - } - - // ========================================================= - // MAIN LOOP (dbuf4) — nested groups x chunks (kept nested; flattening it - // with a conditional coopmat epilog crashes the Xclipse PAL compiler at - // large spec-resolved trip counts). One barrier per chunk. Chunk - // iteration (global index `chunk`): - // 1. barrier — A/B slice (chunk%2) fully written; on the first chunk - // of group g, wsum/wsc slice (g%2) is too. - // 2. prefetch — chunk+1 (A blocks, B blocks) into temp; when chunk+1 - // starts a new group, also its wsum/wsc element. Skipped - // entirely on the final chunk. - // 3. int8 MMA — on slice (chunk%2) into accum_int32. - // 4. store — temp -> A/B slice ((chunk+1)%2), unpacking the weight; - // on a group boundary, wsum/wsc -> slice ((g+1)%2). - // The group epilog runs unconditionally at the tail of each group. - // ========================================================= - uint chunk = 0; - for (uint group_i = 0; group_i < num_groups; ++group_i) { - for (uint inner = 0; inner < CHUNKS_PER_GROUP; ++inner, ++chunk) { - const bool has_next = chunk + 1u < num_chunks; - const bool group_crossing = has_next && (inner + 1u == CHUNKS_PER_GROUP); - const uint cur_a = (chunk % 2u) * ASH_SLICE_U32; - const uint cur_b = (chunk % 2u) * BSH_SLICE_U32; - const uint nxt_a = ((chunk + 1u) % 2u) * ASH_SLICE_U32; - const uint nxt_b = ((chunk + 1u) % 2u) * BSH_SLICE_U32; - - barrier(); - - // --- 2. prefetch chunk+1 -> temp --- - if (has_next) { - const uint chunkK_nxt = (chunk + 1u) * WG_TILE_K; - [[unroll]] for (uint si = 0; si < A_SLOTS_PER_THREAD; ++si) { - const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; - if (slot < A_TOTAL_SLOTS) { - const uint row = slot / K_BLOCKS_PER_CHUNK; - const uint k4 = slot % K_BLOCKS_PER_CHUNK; - temp_A[si] = t_packed_int8_input - [(tile_m_start + row) * nblocks_x_A + (chunkK_nxt >> 2u) + k4]; - } - } -#ifdef WEIGHT_INT4 - [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { - const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; - const uint block_in_chunk = slot >> 3u; - const uint k4_blk = (chunkK_nxt >> 2u) + block_in_chunk / N8_PER_TILE; - const uint n8_blk = (tile_n_start >> 3u) + (block_in_chunk % N8_PER_TILE); -#ifdef WEIGHT_BUFFER - temp_B[si] = t_packed_weight[(n8_blk * nblocks_x_A) + k4_blk]; -#else - temp_B[si] = texelFetch(t_packed_weight, ivec2(k4_blk, n8_blk), 0); -#endif - } - if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { - const uint n_idx = tile_n_start + gl_LocalInvocationID.x; - f16vec4 sv = t_weight_scales[(group_i + 1u) * N4 + (n_idx >> 2u)]; - temp_wsc = float(sv[n_idx & 3u]); - temp_wsum = t_weight_sums[(group_i + 1u) * N + n_idx]; - } -#else - if (b_active) { - const uint block_y_w = (chunkK_nxt >> 2u) + b_k4_in_chunk; - const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; -#ifdef WEIGHT_BUFFER - temp_B = t_packed_weight[(block_y_w * N4) + block_x_w]; -#else - temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, block_y_w), 0); -#endif - } -#endif - } - - // --- 3. int8 MMA on the cur slice --- - [[unroll]] for (uint k = 0; k < NUM_K_SLABS; ++k) { - const uint slab_a_base_u32 = cur_a + k * A_SLAB_U32; - const uint slab_b_base_u32 = cur_b + k * B_SLAB_U32; - - coopmat matA[MMAS_PER_SG_M]; - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - const uint row_a = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); - coopMatLoad( - matA[i], Ash_int8, - slab_a_base_u32 + row_a * A_STRIDE_U32, - A_STRIDE_U32, - gl_CooperativeMatrixLayoutRowMajor); - } - - coopmat matB; - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint col_b = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopMatLoad( - matB, Bsh_int8, - slab_b_base_u32 + col_b * B_STRIDE_U32, - B_STRIDE_U32, - gl_CooperativeMatrixLayoutColumnMajor); - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - accum_int32[i][j] = coopMatMulAdd(matA[i], matB, accum_int32[i][j]); - } - } - } - - // --- 4. store temp (chunk+1) -> nxt slice --- - if (has_next) { - [[unroll]] for (uint si = 0; si < A_SLOTS_PER_THREAD; ++si) { - const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; - if (slot < A_TOTAL_SLOTS) { - const uint row = slot / K_BLOCKS_PER_CHUNK; - const uint k4 = slot % K_BLOCKS_PER_CHUNK; - const uint slab_idx = k4 / (MMA_K >> 2u); - const uint k_uint_in_slab = k4 % (MMA_K >> 2u); - Ash_int8 - [nxt_a + slab_idx * A_SLAB_U32 + row * A_STRIDE_U32 + - k_uint_in_slab] = uint(temp_A[si]); - } - } -#ifdef WEIGHT_INT4 - [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { - const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; - const uint block_in_chunk = slot >> 3u; - const uint col_in_block = slot & 7u; - const uint k4_in_chunk = block_in_chunk / N8_PER_TILE; - const uint n8_in_tile = block_in_chunk % N8_PER_TILE; - const uint r = col_in_block & 3u; - const uint parity = col_in_block >> 2u; - const int w = temp_B[si][r]; - const int base = int(4u * parity); - const int v0 = (((w >> (base + 0)) & 0xF) - 8) & 0xFF; - const int v1 = (((w >> (base + 8)) & 0xF) - 8) & 0xFF; - const int v2 = (((w >> (base + 16)) & 0xF) - 8) & 0xFF; - const int v3 = (((w >> (base + 24)) & 0xF) - 8) & 0xFF; - const uint n_col = n8_in_tile * 8u + r + parity * 4u; - const uint slab_idx = k4_in_chunk / (MMA_K >> 2u); - const uint k4_in_slab = k4_in_chunk % (MMA_K >> 2u); - Bsh_int8[nxt_b + slab_idx * B_SLAB_U32 + n_col * B_STRIDE_U32 + k4_in_slab] = - uint(v0 | (v1 << 8) | (v2 << 16) | (v3 << 24)); - } - if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { - const uint wbase_nxt = ((group_i + 1u) % 2u) * WG_TILE_N; - wsum_sh[wbase_nxt + gl_LocalInvocationID.x] = temp_wsum; - wsc_sh[wbase_nxt + gl_LocalInvocationID.x] = temp_wsc; - } -#else - if (b_active) { - const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); - const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); - const uint n_col_base = b_n_uint_col * 4u; - [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { - Bsh_int8[nxt_b + slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab] = - uint(temp_B[n_in_blk]); - } - } -#endif - } - } // chunks - - // --- Group epilog: dequant accum_int32 -> result, reset accum --- - { - const uint wbase = (group_i % 2u) * WG_TILE_N; - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - - coopmat wsum_bcast; - coopMatLoad( - wsum_bcast, wsum_sh, - wbase + local_n_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutRowMajor); - - coopmat wsc_bcast; - coopMatLoad( - wsc_bcast, wsc_sh, - wbase + local_n_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutRowMajor); - - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - coopmat adjusted = - accum_int32[i][j] - izp_bcast[i] * wsum_bcast; - coopmat adjusted_fp = - coopmat(adjusted); - coopmat scales_outer = - ifs_bcast[i] * wsc_bcast; - result[i][j] += adjusted_fp * scales_outer; - accum_int32[i][j] = coopmat(0); - } - } - } - } // groups - - // --- Bias (optional) --- -#ifdef HAS_BIAS - if (apply_bias > 0) { - for (uint t = gl_LocalInvocationID.x; t < WG_TILE_N; t += WG_SIZE) { - bias_sh[t] = float(t_bias[tile_n_start + t]); - } - memoryBarrierShared(); - barrier(); - } -#endif - - // --- Store result tile --- - // N for the store address math MUST come from the spec constant, not the - // sizes UBO (see out_N_arg above). -#ifdef IO_TEXTURE - // Epilogue iteration i drains accumulator row-block i from EVERY subgroup - // into Csh_out at once, so the SG_GRID_Y bands it holds are disjoint global - // row ranges; the whole workgroup then imageStores them. lr / MMA_M is the - // writing subgroup's warpInTile.y, so the global row reproduces the buffer - // path's gi exactly. - // - // PORTABILITY NOTE: the barrier() in the loop body keeps this loop rolled - // despite [[unroll]], so result[i][j] IS dynamically indexed. Coopmat arrays - // are opaque per-lane storage and dynamic indexing is exactly the construct - // the Xclipse/AMD-PAL compiler has broken before -- check this first if the - // texture variants miscompile on M51. - const uint CSH_TEXELS_PER_ROW = WG_TILE_N / 4u; - const uint CSH_TEXELS = CSH_ROWS * CSH_TEXELS_PER_ROW; - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - // Guards Csh_out against the previous iteration's readers. Inert on i == 0 - // but must stay unconditional to remain workgroup-uniform. - barrier(); - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { -#ifdef HAS_BIAS - if (apply_bias > 0) { - const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopmat bias_tile; - coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); - result[i][j] += bias_tile; - } -#endif - coopmat out_tile = - coopmat(result[i][j]); - coopMatStore( - out_tile, Csh_out, - warpInTile.y * MMA_M * WG_TILE_N + - MMA_N * (MMAS_PER_SG_N * warpInTile.x + j), - WG_TILE_N, - gl_CooperativeMatrixLayoutRowMajor); - } - memoryBarrierShared(); - barrier(); - - for (uint t = gl_LocalInvocationID.x; t < CSH_TEXELS; t += WG_SIZE) { - const uint lr = t / CSH_TEXELS_PER_ROW; - const uint lc4 = t % CSH_TEXELS_PER_ROW; - const uint m = - tile_m_start + (lr / MMA_M) * SG_TILE_M + i * MMA_M + (lr % MMA_M); - const uint base = lr * WG_TILE_N + lc4 * 4u; - imageStore( - t_output, - ivec3(tile_n_start / 4u + lc4, m, 0), - vec4( - float(Csh_out[base]), - float(Csh_out[base + 1u]), - float(Csh_out[base + 2u]), - float(Csh_out[base + 3u]))); - } - } -#else - const uint N_out = uint(out_N_arg); - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint gi = tile_m_start + MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); - const uint gj = tile_n_start + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - -#ifdef HAS_BIAS - if (apply_bias > 0) { - const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopmat bias_tile; - coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); - result[i][j] += bias_tile; - } -#endif - - coopmat out_tile = - coopmat(result[i][j]); - coopMatStore( - out_tile, t_output, - gi * N_out + gj, N_out, - gl_CooperativeMatrixLayoutRowMajor); - } - } -#endif // IO_TEXTURE -} diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trm.yaml b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trm.yaml deleted file mode 100644 index 28f9bdf03e5..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trm.yaml +++ /dev/null @@ -1,188 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -# "-tr" (coopmat-staged A) variants of the int8 dq8ca_q4gsw coopmat kernel, -# ported from shmem_double_buf4-tr.comp (vk_cooperative_matrix_perf, -# gemm-ubm). linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trm.glsl is a fork of -# linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4.glsl with ONLY the A-side global -> -# LDS staging swapped from a per-thread ivec4 copy to -# coopMatLoad(global) -> coopmat<> -> coopMatStore(LDS). B staging is -# byte-identical (int4 nibble unpack cannot be coopmat-staged) and B was -# already column-major in LDS, so it is not a delta. -# -# These variants require t_packed_int8_input in the ROW-MAJOR -# kPackedInt8_4W layout (produced by quantize_and_pack_4w_with_group_sums), -# NOT the stock 4h4w block layout. QuantizedLinear.cpp switches the tensor -# layout and the packer node together with the variant token, and only when -# the coopmat gate passes -- see dq8ca_wants_rowmajor_int8_input(). -# -# Selected via -# ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4trm_txkgs. -# NOT the shipped default: unvalidated until it passes repeated -# test_llama_microbench --correctness-only runs (a single pass is not proof -- -# see dq8ca_coopmat_variant()'s comment). -# -# Tile-geometry preconditions beyond dbuf4's: WG_TILE_M % MMA_M == 0 and -# WG_TILE_K % MMA_K == 0 (both hold for every entry below). The seed set is -# deliberately small -- widen it from a sweep once correctness is established. - -linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trm: - parameter_names_with_default_values: - PRECISION: highp - HAS_BIAS: false - IO_STORAGE: buffer - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - MMA_M: 16 - MMA_N: 16 - MMA_K: 16 - WG_TILE_M: 64 - WG_TILE_N: 32 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - shader_variants: - # production 8da4w tile (dbuf2/dbuf4 default) -- the A/B anchor - # A tiles/chunk = 8, subgroups = 2 -> A_TILES_PER_SG = 4 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trm_t64x32k32g12s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 32 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trm_t64x32k32g12s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 32 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trm_t64x32k32g12s64_buffer_buffer_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: buffer - WG_TILE_M: 64 - WG_TILE_N: 32 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - # q4gsw production geometry; 8 A-tiles over 4 subgroups - # A tiles/chunk = 8, subgroups = 4 -> A_TILES_PER_SG = 2 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trm_t128x128k16g22s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 2 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trm_t128x128k16g22s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 2 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - # 1 A-tile per subgroup (A_TILES_PER_SG == 1, no leftover slot) - # A tiles/chunk = 8, subgroups = 8 -> A_TILES_PER_SG = 1 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trm_t64x64k32g24s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trm_t64x64k32g24s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - # 16 A-tiles over 8 subgroups - # A tiles/chunk = 16, subgroups = 8 -> A_TILES_PER_SG = 2 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trm_t128x64k32g24s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trm_t128x64k32g24s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - # Reference-matching tile: shmem_double_buf4-tr.comp's own geometry - # (TILE 128x128x16, WORKGROUP_WIDTH_IN_SUBGROUPS=4 x DBUF4_GRID_HEIGHT=2, - # wave32) => 8 subgroups / 256 threads, MMAS_PER_SG 4x2. Added so our - # kernel can be A/B'd against the teammate's PAL capture at the same - # tile and the same 2048x1024x4096 shape (8b wk_wv prefill). - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trm_t128x128k16g42s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trm_t128x128k16g42s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - # G2/Pavan's documented INT8 uBM config: TILE 128x128x64 with a 4x2 - # subgroup grid at wave32 (workgroupSize=256). Added to compare our kernel - # against their reference dbuf4 number (~1973us) at the same geometry and - # the same shape (M2048 K4096 N1024 == 8b wk_wv prefill). - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trm_t128x128k64g42s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4trm_t128x128k64g42s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp.glsl b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp.glsl deleted file mode 100644 index 65615a2f1f5..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp.glsl +++ /dev/null @@ -1,687 +0,0 @@ -/* - * Copyright (c) Meta Platforms, Inc. and affiliates. - * All rights reserved. - * - * This source code is licensed under the BSD-style license found in the - * LICENSE file in the root directory of this source tree. - */ - -/* - * "zp-hoisted" variant: identical to dbuf4 except that the activation - * zero-point correction and the per-row activation scale are applied ONCE - * after the group loop instead of once per quantization group. - * - * The per-group epilogue term factors exactly: - * - * out[m][n] = ifs[m] * SUM_g wsc[g][n] * ( acc[g] - izp[m]*wsum[g][n] ) - * = ifs[m] * [ SUM_g wsc[g][n]*acc[g] - izp[m]*SUM_g wsc[g][n]*wsum[g][n] ] - * \__ weight-side only __/ - * - * `ifs` is per-row and group-independent so it factors out entirely, and the - * zero-point term separates into a per-row scalar times a per-output-channel - * weight-side sum. That sum depends on no activation data, so it is - * accumulated once into `wcorr_sh` in the prologue rather than being rebuilt - * per group. - * - * Consequences vs dbuf4, per accumulator tile per group: - * - gone: izp*wsum multiply and the subtract (48 v_sub* in the loop body) - * - gone: ifs*wsc multiply (part of 48 dequant-fp) - * - gone: the wsum_sh shared array and its ping-pong - * - gone: izp_bcast / ifs_bcast live across the loop (register pressure) - * - kept: result += float(acc) * wsc - * - * Exact in exact arithmetic, but NOT bit-exact in fp32 -- the summation order - * changes -- so it is gated on the correctness matrix like any other change. - * - * No new binding and no export-format change: the weight-side sum is derived - * in the prologue from t_weight_scales and t_weight_sums, both already bound. - * - * Selected via - * ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4zp_txkgs<32|64>. - * - * Original dbuf4 header follows. - * - * TILE/SUBGROUP-SWEEP variant of the int8 dq8ca_q4gsw coopmat shader's dbuf4 - * ("store-first-for-next", the ORIGINAL loop structure before specs/025 User - * Story 1 picked dbuf2) loop structure (specs/041-dbuf4-tile-sweep). Forked - * from linear_dq8ca_q4gsw_coopmat_tsweep.glsl (which carries dbuf2's loop, - * the production winner) -- everything except the PROLOGUE/MAIN LOOP block - * is identical: bindings, spec-constants, tile-geometry templating, LDS - * layout (ColumnMajor B + skew), int8 WMMA thread maps, group epilog, - * bias/store epilogue. Only the loop structure is swapped to dbuf4, - * recovered from git commit 8d0f23ee78's - * linear_dq8ca_q4gsw_coopmat_dbuf4.glsl (see specs/041/reference/) -- the - * byte-identical pre-swap copy of what is now linear_dq8ca_qw_coopmat.glsl. - * - * The nested `groups x chunks` loop and unconditional group epilog are kept - * exactly as in dbuf2 -- flattening them crashes the Xclipse PAL compiler at - * large spec-resolved trip counts (see dbuf2's own header). Only the - * store/barrier/prefetch ORDER within each chunk iteration is inverted: - * - * dbuf2 (this file's base): store(temp, already prefetched -> cur slice) - * -> barrier -> MMA(cur) -> prefetch(next -> temp) [store owns the - * CURRENT chunk, at the iteration's start] - * dbuf4 (this file): barrier -> prefetch(next -> temp) -> MMA(cur) -> - * store(temp -> next slice) [store owns the NEXT chunk, at the - * iteration's end -- the mirror image] - * - * The group wsum/wsc ping-pong is inverted the same way: dbuf4 stores the - * next group's values (prefetched during the crossing chunk) at the TAIL of - * that chunk, instead of dbuf2's HEAD-of-new-group placement. - * - * Selected at dispatch via - * ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4_txkgs<32|64> - * (QuantizedLinear.cpp), additive to the existing tsweep_t... (dbuf2) - * namespace. - * - * KHR Cooperative Matrix variant of the dynamically-quantized-activation - * linear tiled shader (WEIGHT_NBITS=4): - * 4 -> linear_dq8ca_q4gsw_coopmat INT4 group-symmetric weight - * - * Performs: out[M,N] = dequant(int8_act) * dequant(int_w) (+ bias) - * via coopmat x coopmat -> coopmat on the matrix unit. - * - * Hard preconditions: - * M % WG_TILE_M == 0, N % WG_TILE_N == 0, K % WG_TILE_K == 0, - * INT4: group_size % WG_TILE_K == 0, - * device exposes coopmatx-> at 16x16x16. - */ - -#version 450 core - -#extension GL_KHR_cooperative_matrix : require -#extension GL_KHR_memory_scope_semantics : require -#extension GL_KHR_shader_subgroup_basic : enable -#extension GL_EXT_shader_explicit_arithmetic_types : require -#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require -#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require -#extension GL_EXT_control_flow_attributes : enable - -#define PRECISION ${PRECISION} - -$if WEIGHT_NBITS == 4: - #define WEIGHT_INT4 - -$if HAS_BIAS: - #define HAS_BIAS - -$if WEIGHT_STORAGE == "buffer": - #define WEIGHT_BUFFER - -$if IO_STORAGE == "texture3d": - #define IO_TEXTURE - -layout(std430) buffer; - -#include "common.glslh" - -// Bindings — match add_linear_dqa_qw_node arg order: -// output(0), fp_input(1), packed_int8_input(2), int_input_sums(3 - unused), -// input_scales(4), input_zps(5), packed_weight(6), weight_sums(7), -// weight_scales(8), bias(9). -${layout_declare_tensor(B, "w", "t_output", "half", IO_STORAGE, is_scalar_array=True)} -// t_input is unread here -- the activations arrive already quantized in -// t_packed_int8_input -- but stays declared so the binding layout matches the -// dispatch site. It tracks IO_STORAGE so the two IO tensors stay consistent. -${layout_declare_tensor(B, "r", "t_input", "half", IO_STORAGE, is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_packed_int8_input", "int", "buffer", is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_int8_input_sums", "int", "buffer", is_scalar_array=True)} -${layout_declare_tensor(B, "r", "t_int8_input_scales", "half", "texture3d")} -${layout_declare_tensor(B, "r", "t_int8_input_zps", "int8", "texture3d")} -${layout_declare_tensor(B, "r", "t_packed_weight", "int", WEIGHT_STORAGE, is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_weight_sums", "int", "buffer", is_scalar_array=True)} -${layout_declare_tensor(B, "r", "t_weight_scales", "half", "buffer", is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_bias", "half", "buffer", is_scalar_array=True)} - -${layout_declare_ubo(B, "ivec4", "output_sizes")} -${layout_declare_ubo(B, "ivec4", "input_sizes")} - -layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; - -${layout_declare_spec_const(C, "int", "apply_bias", "0")} -// INT4 only; inert (0) for INT8 so the dispatcher's spec list lines up. -${layout_declare_spec_const(C, "int", "K4_per_group", "0")} -${layout_declare_spec_const(C, "int", "num_groups_arg", "0")} -${layout_declare_spec_const(C, "int", "out_N_arg", "0")} - -// Tile geometry -const uint MMA_M = ${MMA_M}; -const uint MMA_N = ${MMA_N}; -const uint MMA_K = ${MMA_K}; - -const uint WG_TILE_M = ${WG_TILE_M}; -const uint WG_TILE_N = ${WG_TILE_N}; -const uint WG_TILE_K = ${WG_TILE_K}; - -const uint SG_GRID_X = ${SG_GRID_X}; -const uint SG_GRID_Y = ${SG_GRID_Y}; -const uint SUBGROUP_SIZE = ${SUBGROUP_SIZE}; -const uint NUM_SUBGROUPS = SG_GRID_X * SG_GRID_Y; -const uint WG_SIZE = NUM_SUBGROUPS * SUBGROUP_SIZE; - -const uint SG_TILE_M = WG_TILE_M / SG_GRID_Y; -const uint SG_TILE_N = WG_TILE_N / SG_GRID_X; -const uint MMAS_PER_SG_M = SG_TILE_M / MMA_M; -const uint MMAS_PER_SG_N = SG_TILE_N / MMA_N; - -const uint A_SLAB_INT8 = WG_TILE_M * MMA_K; -const uint B_USEFUL_U32 = MMA_K / 4u; -const uint B_STRIDE_U32 = B_USEFUL_U32 + 1u; // +1 skew -const uint B_SLAB_U32 = WG_TILE_N * B_STRIDE_U32; -const uint NUM_K_SLABS = WG_TILE_K / MMA_K; - -const uint A_SLAB_U32 = A_SLAB_INT8 / 4u; -const uint A_STRIDE_U32 = MMA_K / 4u; - -// One ping-pong slice covers all K-slabs of one chunk. -const uint ASH_SLICE_U32 = NUM_K_SLABS * A_SLAB_U32; -const uint BSH_SLICE_U32 = NUM_K_SLABS * B_SLAB_U32; - -// Double-buffered MMA operand staging. -shared uint Ash_int8[2u * ASH_SLICE_U32]; -shared uint Bsh_int8[2u * BSH_SLICE_U32]; - -// Per-WG-tile-row activation params (loaded ONCE at WG start; constant -// across groups). -shared int izp_sh[WG_TILE_M]; // int32 (cast from int8 source) for broadcast -shared float ifs_sh[WG_TILE_M]; // float32 (cast from fp16 source) for broadcast - -// Per-(group, output-channel) weight params, ping-ponged by group parity. -// (For per-channel INT8 only slice 0 is ever used.) -shared float wsc_sh[2u * WG_TILE_N]; -// SUM_g wsc[g][n]*wsum[g][n] per output channel -- weight-side only, so it is -// accumulated once in the prologue. Replaces dbuf4's ping-ponged wsum_sh -// (which was 2*WG_TILE_N ints), so this is a net LDS saving. -shared float wcorr_sh[WG_TILE_N]; - -#ifdef HAS_BIAS -shared float bias_sh[WG_TILE_N]; -#endif - -#ifdef IO_TEXTURE -// Result staging for the imageStore epilogue, mirroring the fp16 kernel: -// SG_GRID_Y bands of MMA_M rows, each WG_TILE_N wide, row-major. A full -// WG_TILE_M x WG_TILE_N buffer would cost SG_GRID_Y/MMAS_PER_SG_M x more LDS -// and wreck occupancy. float16_t-typed because coopMatStore needs it. -const uint CSH_ROWS = SG_GRID_Y * MMA_M; -shared float16_t Csh_out[CSH_ROWS * WG_TILE_N]; -#endif - -// Running fp32 accumulator (across all groups). -coopmat - result[MMAS_PER_SG_M][MMAS_PER_SG_N]; - -// Per-group int32 MMA accumulator. -coopmat - accum_int32[MMAS_PER_SG_M][MMAS_PER_SG_N]; - -void main() { - const uvec2 tileID = uvec2(gl_WorkGroupID.xy); - const uvec2 warpInTile = uvec2( - gl_SubgroupID % SG_GRID_X, - gl_SubgroupID / SG_GRID_X); - - const uint K = uint(input_sizes.x); - const uint N = uint(output_sizes.x); - const uint N4 = (N + 3u) / 4u; - const uint nblocks_x_A = (K + 3u) >> 2u; - -#ifdef WEIGHT_INT4 - const uint num_groups = uint(num_groups_arg); - const uint CHUNKS_PER_GROUP = uint(K4_per_group) * 4u / WG_TILE_K; -#else - // Per-channel: a single quant "group" spanning all of K. The nested - // groups x chunks loop below collapses to a flat chunk loop, the wsum/wsc - // ping-pong never crosses a boundary, and the epilog runs exactly once. - const uint num_groups = 1u; - const uint CHUNKS_PER_GROUP = uint(num_groups_arg); -#endif - const uint num_chunks = num_groups * CHUNKS_PER_GROUP; - - const uint tile_m_start = WG_TILE_M * tileID.y; - const uint tile_n_start = WG_TILE_N * tileID.x; - - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - result[i][j] = coopmat(0.0); - accum_int32[i][j] = coopmat(0); - } - } - - // --- A staging thread map: one (m4, k4) ivec4 block per active thread --- - const uint K_BLOCKS_PER_CHUNK = WG_TILE_K >> 2u; - const uint A_ACTIVE_THREADS = (WG_TILE_M >> 2u) * K_BLOCKS_PER_CHUNK; - const uint a_m_block = gl_LocalInvocationID.x / K_BLOCKS_PER_CHUNK; - const uint a_k_block = gl_LocalInvocationID.x % K_BLOCKS_PER_CHUNK; - const bool a_active = gl_LocalInvocationID.x < A_ACTIVE_THREADS; - -#ifdef WEIGHT_INT4 - // --- B staging thread map: (block, col) slots; each slot extracts one - // ColumnMajor LDS uint (4 K-contiguous sign-extended int8) --- - const uint B_TOTAL_SLOTS = K_BLOCKS_PER_CHUNK * WG_TILE_N; - const uint B_SLOTS_PER_THREAD = B_TOTAL_SLOTS / WG_SIZE; - const uint N8_PER_TILE = WG_TILE_N >> 3u; -#else - // --- B staging thread map: one (k4, n4) ivec4 block per active thread --- - const uint B_FETCH_SLOTS = K_BLOCKS_PER_CHUNK * (WG_TILE_N >> 2u); - const uint N4_PER_TILE = WG_TILE_N >> 2u; - const uint b_k4_in_chunk = gl_LocalInvocationID.x / N4_PER_TILE; - const uint b_n_uint_col = gl_LocalInvocationID.x % N4_PER_TILE; - const bool b_active = gl_LocalInvocationID.x < B_FETCH_SLOTS; -#endif - - // Prefetch temp registers. - ivec4 temp_A; -#ifdef WEIGHT_INT4 - ivec4 temp_B[B_SLOTS_PER_THREAD]; - float temp_wsc; -#else - ivec4 temp_B; -#endif - - // ========================================================= - // PROLOGUE - // ========================================================= - if (gl_LocalInvocationID.x < (WG_TILE_M >> 2u)) { - const uint m4 = (tile_m_start >> 2u) + gl_LocalInvocationID.x; - const vec4 sc = vec4(texelFetch(t_int8_input_scales, ivec3(m4, 0, 0), 0)); - const ivec4 zp = texelFetch(t_int8_input_zps, ivec3(m4, 0, 0), 0); - const uint base = gl_LocalInvocationID.x * 4u; - ifs_sh[base + 0u] = sc.x; ifs_sh[base + 1u] = sc.y; - ifs_sh[base + 2u] = sc.z; ifs_sh[base + 3u] = sc.w; - izp_sh[base + 0u] = zp.x; izp_sh[base + 1u] = zp.y; - izp_sh[base + 2u] = zp.z; izp_sh[base + 3u] = zp.w; - } - // Group 0 weight scales -> slice 0, and the hoisted weight-side correction - // SUM_g wsc[g][n]*wsum[g][n] accumulated across ALL groups. The loop is - // prologue-only (the prologue is ~1.4% of the dynamic instruction stream), - // and it replaces per-group wsum work inside the loop body. - if (gl_LocalInvocationID.x < WG_TILE_N) { - const uint n_idx = tile_n_start + gl_LocalInvocationID.x; - f16vec4 sv0 = t_weight_scales[n_idx >> 2u]; - wsc_sh[gl_LocalInvocationID.x] = float(sv0[n_idx & 3u]); - - float corr = 0.0; - for (uint g = 0; g < num_groups; ++g) { - f16vec4 sv = t_weight_scales[g * N4 + (n_idx >> 2u)]; - corr += float(sv[n_idx & 3u]) * float(t_weight_sums[g * N + n_idx]); - } - wcorr_sh[gl_LocalInvocationID.x] = corr; - } - memoryBarrierShared(); - barrier(); - - // NOTE: dbuf4 builds izp_bcast/ifs_bcast here and keeps them live across the - // whole group loop. This variant needs them only AFTER the loop, so they are - // loaded there instead -- that is the register-pressure saving. - - // dbuf4: prefetch chunk 0 into temp registers, THEN store to slice 0 (no - // barrier here -- the main loop's first iteration barriers before - // reading slice 0). - if (a_active) { - const uint m4_global = (tile_m_start >> 2u) + a_m_block; - temp_A = t_packed_int8_input[m4_global * nblocks_x_A + a_k_block]; - } -#ifdef WEIGHT_INT4 - [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { - const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; - const uint block_in_chunk = slot >> 3u; - const uint k4_blk = block_in_chunk / N8_PER_TILE; - const uint n8_blk = (tile_n_start >> 3u) + (block_in_chunk % N8_PER_TILE); -#ifdef WEIGHT_BUFFER - temp_B[si] = t_packed_weight[(n8_blk * nblocks_x_A) + k4_blk]; -#else - temp_B[si] = texelFetch(t_packed_weight, ivec2(k4_blk, n8_blk), 0); -#endif - } -#else - if (b_active) { - const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; -#ifdef WEIGHT_BUFFER - temp_B = t_packed_weight[(b_k4_in_chunk * N4) + block_x_w]; -#else - temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, b_k4_in_chunk), 0); -#endif - } -#endif - { - // store chunk 0 -> slice 0 - if (a_active) { - const uint slab_idx = a_k_block / (MMA_K >> 2u); - const uint k_uint_in_slab = a_k_block % (MMA_K >> 2u); - const uint base_row = a_m_block * 4u; - [[unroll]] for (uint m4i = 0; m4i < 4u; ++m4i) { - Ash_int8[slab_idx * A_SLAB_U32 + (base_row + m4i) * A_STRIDE_U32 + k_uint_in_slab] = - uint(temp_A[m4i]); - } - } -#ifdef WEIGHT_INT4 - [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { - const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; - const uint block_in_chunk = slot >> 3u; - const uint col_in_block = slot & 7u; - const uint k4_in_chunk = block_in_chunk / N8_PER_TILE; - const uint n8_in_tile = block_in_chunk % N8_PER_TILE; - const uint r = col_in_block & 3u; - const uint parity = col_in_block >> 2u; - const int w = temp_B[si][r]; - const int base = int(4u * parity); - const int v0 = (((w >> (base + 0)) & 0xF) - 8) & 0xFF; - const int v1 = (((w >> (base + 8)) & 0xF) - 8) & 0xFF; - const int v2 = (((w >> (base + 16)) & 0xF) - 8) & 0xFF; - const int v3 = (((w >> (base + 24)) & 0xF) - 8) & 0xFF; - const uint n_col = n8_in_tile * 8u + r + parity * 4u; - const uint slab_idx = k4_in_chunk / (MMA_K >> 2u); - const uint k4_in_slab = k4_in_chunk % (MMA_K >> 2u); - Bsh_int8[slab_idx * B_SLAB_U32 + n_col * B_STRIDE_U32 + k4_in_slab] = - uint(v0 | (v1 << 8) | (v2 << 16) | (v3 << 24)); - } -#else - if (b_active) { - const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); - const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); - const uint n_col_base = b_n_uint_col * 4u; - [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { - Bsh_int8[slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab] = - uint(temp_B[n_in_blk]); - } - } -#endif - } - - // ========================================================= - // MAIN LOOP (dbuf4) — nested groups x chunks (kept nested; flattening it - // with a conditional coopmat epilog crashes the Xclipse PAL compiler at - // large spec-resolved trip counts). One barrier per chunk. Chunk - // iteration (global index `chunk`): - // 1. barrier — A/B slice (chunk%2) fully written; on the first chunk - // of group g, wsum/wsc slice (g%2) is too. - // 2. prefetch — chunk+1 (A blocks, B blocks) into temp; when chunk+1 - // starts a new group, also its wsum/wsc element. Skipped - // entirely on the final chunk. - // 3. int8 MMA — on slice (chunk%2) into accum_int32. - // 4. store — temp -> A/B slice ((chunk+1)%2), unpacking the weight; - // on a group boundary, wsum/wsc -> slice ((g+1)%2). - // The group epilog runs unconditionally at the tail of each group. - // ========================================================= - uint chunk = 0; - for (uint group_i = 0; group_i < num_groups; ++group_i) { - for (uint inner = 0; inner < CHUNKS_PER_GROUP; ++inner, ++chunk) { - const bool has_next = chunk + 1u < num_chunks; - const bool group_crossing = has_next && (inner + 1u == CHUNKS_PER_GROUP); - const uint cur_a = (chunk % 2u) * ASH_SLICE_U32; - const uint cur_b = (chunk % 2u) * BSH_SLICE_U32; - const uint nxt_a = ((chunk + 1u) % 2u) * ASH_SLICE_U32; - const uint nxt_b = ((chunk + 1u) % 2u) * BSH_SLICE_U32; - - barrier(); - - // --- 2. prefetch chunk+1 -> temp --- - if (has_next) { - const uint chunkK_nxt = (chunk + 1u) * WG_TILE_K; - if (a_active) { - const uint m4_global = (tile_m_start >> 2u) + a_m_block; - const uint k4_global = (chunkK_nxt >> 2u) + a_k_block; - temp_A = t_packed_int8_input[m4_global * nblocks_x_A + k4_global]; - } -#ifdef WEIGHT_INT4 - [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { - const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; - const uint block_in_chunk = slot >> 3u; - const uint k4_blk = (chunkK_nxt >> 2u) + block_in_chunk / N8_PER_TILE; - const uint n8_blk = (tile_n_start >> 3u) + (block_in_chunk % N8_PER_TILE); -#ifdef WEIGHT_BUFFER - temp_B[si] = t_packed_weight[(n8_blk * nblocks_x_A) + k4_blk]; -#else - temp_B[si] = texelFetch(t_packed_weight, ivec2(k4_blk, n8_blk), 0); -#endif - } - if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { - const uint n_idx = tile_n_start + gl_LocalInvocationID.x; - f16vec4 sv = t_weight_scales[(group_i + 1u) * N4 + (n_idx >> 2u)]; - temp_wsc = float(sv[n_idx & 3u]); - } -#else - if (b_active) { - const uint block_y_w = (chunkK_nxt >> 2u) + b_k4_in_chunk; - const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; -#ifdef WEIGHT_BUFFER - temp_B = t_packed_weight[(block_y_w * N4) + block_x_w]; -#else - temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, block_y_w), 0); -#endif - } -#endif - } - - // --- 3. int8 MMA on the cur slice --- - [[unroll]] for (uint k = 0; k < NUM_K_SLABS; ++k) { - const uint slab_a_base_u32 = cur_a + k * A_SLAB_U32; - const uint slab_b_base_u32 = cur_b + k * B_SLAB_U32; - - coopmat matA[MMAS_PER_SG_M]; - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - const uint row_a = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); - coopMatLoad( - matA[i], Ash_int8, - slab_a_base_u32 + row_a * A_STRIDE_U32, - A_STRIDE_U32, - gl_CooperativeMatrixLayoutRowMajor); - } - - coopmat matB; - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint col_b = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopMatLoad( - matB, Bsh_int8, - slab_b_base_u32 + col_b * B_STRIDE_U32, - B_STRIDE_U32, - gl_CooperativeMatrixLayoutColumnMajor); - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - accum_int32[i][j] = coopMatMulAdd(matA[i], matB, accum_int32[i][j]); - } - } - } - - // --- 4. store temp (chunk+1) -> nxt slice --- - if (has_next) { - if (a_active) { - const uint slab_idx = a_k_block / (MMA_K >> 2u); - const uint k_uint_in_slab = a_k_block % (MMA_K >> 2u); - const uint base_row = a_m_block * 4u; - [[unroll]] for (uint m4i = 0; m4i < 4u; ++m4i) { - Ash_int8[nxt_a + slab_idx * A_SLAB_U32 + (base_row + m4i) * A_STRIDE_U32 + k_uint_in_slab] = - uint(temp_A[m4i]); - } - } -#ifdef WEIGHT_INT4 - [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { - const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; - const uint block_in_chunk = slot >> 3u; - const uint col_in_block = slot & 7u; - const uint k4_in_chunk = block_in_chunk / N8_PER_TILE; - const uint n8_in_tile = block_in_chunk % N8_PER_TILE; - const uint r = col_in_block & 3u; - const uint parity = col_in_block >> 2u; - const int w = temp_B[si][r]; - const int base = int(4u * parity); - const int v0 = (((w >> (base + 0)) & 0xF) - 8) & 0xFF; - const int v1 = (((w >> (base + 8)) & 0xF) - 8) & 0xFF; - const int v2 = (((w >> (base + 16)) & 0xF) - 8) & 0xFF; - const int v3 = (((w >> (base + 24)) & 0xF) - 8) & 0xFF; - const uint n_col = n8_in_tile * 8u + r + parity * 4u; - const uint slab_idx = k4_in_chunk / (MMA_K >> 2u); - const uint k4_in_slab = k4_in_chunk % (MMA_K >> 2u); - Bsh_int8[nxt_b + slab_idx * B_SLAB_U32 + n_col * B_STRIDE_U32 + k4_in_slab] = - uint(v0 | (v1 << 8) | (v2 << 16) | (v3 << 24)); - } - if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { - const uint wbase_nxt = ((group_i + 1u) % 2u) * WG_TILE_N; - wsc_sh[wbase_nxt + gl_LocalInvocationID.x] = temp_wsc; - } -#else - if (b_active) { - const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); - const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); - const uint n_col_base = b_n_uint_col * 4u; - [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { - Bsh_int8[nxt_b + slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab] = - uint(temp_B[n_in_blk]); - } - } -#endif - } - } // chunks - - // --- Group epilog: scale-only accumulate, reset accum --- - // Just result += float(acc) * wsc. The zero-point subtract and the ifs - // multiply are hoisted out of the group loop (applied once below). - { - const uint wbase = (group_i % 2u) * WG_TILE_N; - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - - coopmat wsc_bcast; - coopMatLoad( - wsc_bcast, wsc_sh, - wbase + local_n_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutRowMajor); - - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - result[i][j] += - coopmat( - accum_int32[i][j]) * wsc_bcast; - accum_int32[i][j] = coopmat(0); - } - } - } - } // groups - - // --- Hoisted correction, applied ONCE: --------------------------------- - // result = ifs * ( result - izp * SUM_g wsc_g*wsum_g ) - // izp/ifs are loaded here rather than before the group loop so they are not - // live across it. - { - coopmat - izpf_bcast[MMAS_PER_SG_M]; - coopmat - ifs_bcast[MMAS_PER_SG_M]; - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - const uint local_m_base = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); - coopmat izp_i; - coopMatLoad( - izp_i, izp_sh, local_m_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutColumnMajor); - izpf_bcast[i] = - coopmat(izp_i); - coopMatLoad( - ifs_bcast[i], ifs_sh, local_m_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutColumnMajor); - } - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopmat wcorr_bcast; - coopMatLoad( - wcorr_bcast, wcorr_sh, local_n_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutRowMajor); - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - result[i][j] = ifs_bcast[i] * (result[i][j] - izpf_bcast[i] * wcorr_bcast); - } - } - } - - // --- Bias (optional) --- -#ifdef HAS_BIAS - if (apply_bias > 0) { - for (uint t = gl_LocalInvocationID.x; t < WG_TILE_N; t += WG_SIZE) { - bias_sh[t] = float(t_bias[tile_n_start + t]); - } - memoryBarrierShared(); - barrier(); - } -#endif - - // --- Store result tile --- - // N for the store address math MUST come from the spec constant, not the - // sizes UBO (see out_N_arg above). -#ifdef IO_TEXTURE - // Epilogue iteration i drains accumulator row-block i from EVERY subgroup - // into Csh_out at once, so the SG_GRID_Y bands it holds are disjoint global - // row ranges; the whole workgroup then imageStores them. lr / MMA_M is the - // writing subgroup's warpInTile.y, so the global row reproduces the buffer - // path's gi exactly. - // - // PORTABILITY NOTE: the barrier() in the loop body keeps this loop rolled - // despite [[unroll]], so result[i][j] IS dynamically indexed. Coopmat arrays - // are opaque per-lane storage and dynamic indexing is exactly the construct - // the Xclipse/AMD-PAL compiler has broken before -- check this first if the - // texture variants miscompile on M51. - const uint CSH_TEXELS_PER_ROW = WG_TILE_N / 4u; - const uint CSH_TEXELS = CSH_ROWS * CSH_TEXELS_PER_ROW; - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - // Guards Csh_out against the previous iteration's readers. Inert on i == 0 - // but must stay unconditional to remain workgroup-uniform. - barrier(); - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { -#ifdef HAS_BIAS - if (apply_bias > 0) { - const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopmat bias_tile; - coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); - result[i][j] += bias_tile; - } -#endif - coopmat out_tile = - coopmat(result[i][j]); - coopMatStore( - out_tile, Csh_out, - warpInTile.y * MMA_M * WG_TILE_N + - MMA_N * (MMAS_PER_SG_N * warpInTile.x + j), - WG_TILE_N, - gl_CooperativeMatrixLayoutRowMajor); - } - memoryBarrierShared(); - barrier(); - - for (uint t = gl_LocalInvocationID.x; t < CSH_TEXELS; t += WG_SIZE) { - const uint lr = t / CSH_TEXELS_PER_ROW; - const uint lc4 = t % CSH_TEXELS_PER_ROW; - const uint m = - tile_m_start + (lr / MMA_M) * SG_TILE_M + i * MMA_M + (lr % MMA_M); - const uint base = lr * WG_TILE_N + lc4 * 4u; - imageStore( - t_output, - ivec3(tile_n_start / 4u + lc4, m, 0), - vec4( - float(Csh_out[base]), - float(Csh_out[base + 1u]), - float(Csh_out[base + 2u]), - float(Csh_out[base + 3u]))); - } - } -#else - const uint N_out = uint(out_N_arg); - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint gi = tile_m_start + MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); - const uint gj = tile_n_start + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - -#ifdef HAS_BIAS - if (apply_bias > 0) { - const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopmat bias_tile; - coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); - result[i][j] += bias_tile; - } -#endif - - coopmat out_tile = - coopmat(result[i][j]); - coopMatStore( - out_tile, t_output, - gi * N_out + gj, N_out, - gl_CooperativeMatrixLayoutRowMajor); - } - } -#endif // IO_TEXTURE -} diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp.yaml b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp.yaml deleted file mode 100644 index 40e7a12f29d..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp.yaml +++ /dev/null @@ -1,187 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -# "zp-hoisted" variants of the int8 dq8ca_q4gsw coopmat kernel: identical to -# tsweep_dbuf4 except the activation zero-point correction and the per-row -# activation scale are applied once after the group loop instead of once per -# quantization group. No new binding, no export-format change -- the weight-side -# correction sum is derived in the prologue from tensors already bound. -# -# Selected via -# ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4zp_txkgs. -# NOT the shipped default until it passes repeated correctness runs. - -linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp: - parameter_names_with_default_values: - PRECISION: highp - HAS_BIAS: false - IO_STORAGE: buffer - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - MMA_M: 16 - MMA_N: 16 - MMA_K: 16 - WG_TILE_M: 64 - WG_TILE_N: 32 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - shader_variants: - # production tile -- the A/B anchor against the 1.3 baseline - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp_t64x32k32g12s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 32 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp_t64x32k32g12s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 32 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp_t64x32k32g12s64_buffer_buffer_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: buffer - WG_TILE_M: 64 - WG_TILE_N: 32 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - # NOTE: t128x128k64g42s32 was seeded here and then REMOVED. Its A-staging - # thread map cannot cover its blocks: A_ACTIVE_THREADS = (128/4)*(64/4) = - # 512 but WG_SIZE = 4*2*32 = 256, so half of A is never staged and the - # kernel is numerically wrong (measured 12/12 correctness failures on the - # dbuf4 equivalent). Use g44s32 for that tile instead -- 4*4*32 = 512 = A. - # deep K + big tile; the reference's winning geometry, A map valid at 512=512 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp_t128x128k64g44s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 4 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp_t128x128k64g44s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 4 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - # task 5.1's explicit ask - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp_t128x128k32g42s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 32 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp_t128x128k32g42s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 32 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - # same tile, more subgroups -> MMAS 2x2 instead of 4x2 (less accumulator pressure) - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp_t128x128k32g44s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 32 - SG_GRID_X: 4 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp_t128x128k32g44s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 32 - SG_GRID_X: 4 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - # moderate step up from the shipped 64x32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp_t128x64k32g42s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp_t128x64k32g42s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - # small step up - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp_t64x64k32g22s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 2 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp_t64x64k32g22s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 2 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - # big tile, shallow K -- isolates tile area from K depth - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp_t128x128k16g42s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zp_t128x128k16g42s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpb.glsl b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpb.glsl deleted file mode 100644 index 1b90a3bb647..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpb.glsl +++ /dev/null @@ -1,714 +0,0 @@ -/* - * Copyright (c) Meta Platforms, Inc. and affiliates. - * All rights reserved. - * - * This source code is licensed under the BSD-style license found in the - * LICENSE file in the root directory of this source tree. - */ - -/* - * "zp-hoisted" variant: identical to dbuf4 except that the activation - * zero-point correction and the per-row activation scale are applied ONCE - * after the group loop instead of once per quantization group. - * - * The per-group epilogue term factors exactly: - * - * out[m][n] = ifs[m] * SUM_g wsc[g][n] * ( acc[g] - izp[m]*wsum[g][n] ) - * = ifs[m] * [ SUM_g wsc[g][n]*acc[g] - izp[m]*SUM_g wsc[g][n]*wsum[g][n] ] - * \__ weight-side only __/ - * - * `ifs` is per-row and group-independent so it factors out entirely, and the - * zero-point term separates into a per-row scalar times a per-output-channel - * weight-side sum. That sum depends on no activation data, so it is - * accumulated once into `wcorr_sh` in the prologue rather than being rebuilt - * per group. - * - * Consequences vs dbuf4, per accumulator tile per group: - * - gone: izp*wsum multiply and the subtract (48 v_sub* in the loop body) - * - gone: ifs*wsc multiply (part of 48 dequant-fp) - * - gone: the wsum_sh shared array and its ping-pong - * - gone: izp_bcast / ifs_bcast live across the loop (register pressure) - * - kept: result += float(acc) * wsc - * - * Exact in exact arithmetic, but NOT bit-exact in fp32 -- the summation order - * changes -- so it is gated on the correctness matrix like any other change. - * - * No new binding and no export-format change: the weight-side sum is derived - * in the prologue from t_weight_scales and t_weight_sums, both already bound. - * - * Additionally widens int4 -> int8 byte-parallel (see widen_nibbles below), - * replacing the per-nibble shift/mask/bias-subtract chain. Bit-identical. - * - * Additionally templates the B LDS skew (B_SKEW) so a power-of-two stride can - * be measured against the baseline +1 skew. - * - * Selected via - * ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4zpb_txkgs<32|64>. - * - * Original dbuf4 header follows. - * - * TILE/SUBGROUP-SWEEP variant of the int8 dq8ca_q4gsw coopmat shader's dbuf4 - * ("store-first-for-next", the ORIGINAL loop structure before specs/025 User - * Story 1 picked dbuf2) loop structure (specs/041-dbuf4-tile-sweep). Forked - * from linear_dq8ca_q4gsw_coopmat_tsweep.glsl (which carries dbuf2's loop, - * the production winner) -- everything except the PROLOGUE/MAIN LOOP block - * is identical: bindings, spec-constants, tile-geometry templating, LDS - * layout (ColumnMajor B + skew), int8 WMMA thread maps, group epilog, - * bias/store epilogue. Only the loop structure is swapped to dbuf4, - * recovered from git commit 8d0f23ee78's - * linear_dq8ca_q4gsw_coopmat_dbuf4.glsl (see specs/041/reference/) -- the - * byte-identical pre-swap copy of what is now linear_dq8ca_qw_coopmat.glsl. - * - * The nested `groups x chunks` loop and unconditional group epilog are kept - * exactly as in dbuf2 -- flattening them crashes the Xclipse PAL compiler at - * large spec-resolved trip counts (see dbuf2's own header). Only the - * store/barrier/prefetch ORDER within each chunk iteration is inverted: - * - * dbuf2 (this file's base): store(temp, already prefetched -> cur slice) - * -> barrier -> MMA(cur) -> prefetch(next -> temp) [store owns the - * CURRENT chunk, at the iteration's start] - * dbuf4 (this file): barrier -> prefetch(next -> temp) -> MMA(cur) -> - * store(temp -> next slice) [store owns the NEXT chunk, at the - * iteration's end -- the mirror image] - * - * The group wsum/wsc ping-pong is inverted the same way: dbuf4 stores the - * next group's values (prefetched during the crossing chunk) at the TAIL of - * that chunk, instead of dbuf2's HEAD-of-new-group placement. - * - * Selected at dispatch via - * ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4_txkgs<32|64> - * (QuantizedLinear.cpp), additive to the existing tsweep_t... (dbuf2) - * namespace. - * - * KHR Cooperative Matrix variant of the dynamically-quantized-activation - * linear tiled shader (WEIGHT_NBITS=4): - * 4 -> linear_dq8ca_q4gsw_coopmat INT4 group-symmetric weight - * - * Performs: out[M,N] = dequant(int8_act) * dequant(int_w) (+ bias) - * via coopmat x coopmat -> coopmat on the matrix unit. - * - * Hard preconditions: - * M % WG_TILE_M == 0, N % WG_TILE_N == 0, K % WG_TILE_K == 0, - * INT4: group_size % WG_TILE_K == 0, - * device exposes coopmatx-> at 16x16x16. - */ - -#version 450 core - -#extension GL_KHR_cooperative_matrix : require -#extension GL_KHR_memory_scope_semantics : require -#extension GL_KHR_shader_subgroup_basic : enable -#extension GL_EXT_shader_explicit_arithmetic_types : require -#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require -#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require -#extension GL_EXT_control_flow_attributes : enable - -#define PRECISION ${PRECISION} - -$if WEIGHT_NBITS == 4: - #define WEIGHT_INT4 - -$if HAS_BIAS: - #define HAS_BIAS - -$if WEIGHT_STORAGE == "buffer": - #define WEIGHT_BUFFER - -$if IO_STORAGE == "texture3d": - #define IO_TEXTURE - -layout(std430) buffer; - -#include "common.glslh" - -// Bindings — match add_linear_dqa_qw_node arg order: -// output(0), fp_input(1), packed_int8_input(2), int_input_sums(3 - unused), -// input_scales(4), input_zps(5), packed_weight(6), weight_sums(7), -// weight_scales(8), bias(9). -${layout_declare_tensor(B, "w", "t_output", "half", IO_STORAGE, is_scalar_array=True)} -// t_input is unread here -- the activations arrive already quantized in -// t_packed_int8_input -- but stays declared so the binding layout matches the -// dispatch site. It tracks IO_STORAGE so the two IO tensors stay consistent. -${layout_declare_tensor(B, "r", "t_input", "half", IO_STORAGE, is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_packed_int8_input", "int", "buffer", is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_int8_input_sums", "int", "buffer", is_scalar_array=True)} -${layout_declare_tensor(B, "r", "t_int8_input_scales", "half", "texture3d")} -${layout_declare_tensor(B, "r", "t_int8_input_zps", "int8", "texture3d")} -${layout_declare_tensor(B, "r", "t_packed_weight", "int", WEIGHT_STORAGE, is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_weight_sums", "int", "buffer", is_scalar_array=True)} -${layout_declare_tensor(B, "r", "t_weight_scales", "half", "buffer", is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_bias", "half", "buffer", is_scalar_array=True)} - -${layout_declare_ubo(B, "ivec4", "output_sizes")} -${layout_declare_ubo(B, "ivec4", "input_sizes")} - -layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; - -${layout_declare_spec_const(C, "int", "apply_bias", "0")} -// INT4 only; inert (0) for INT8 so the dispatcher's spec list lines up. -${layout_declare_spec_const(C, "int", "K4_per_group", "0")} -${layout_declare_spec_const(C, "int", "num_groups_arg", "0")} -${layout_declare_spec_const(C, "int", "out_N_arg", "0")} - -// Tile geometry -const uint MMA_M = ${MMA_M}; -const uint MMA_N = ${MMA_N}; -const uint MMA_K = ${MMA_K}; - -const uint WG_TILE_M = ${WG_TILE_M}; -const uint WG_TILE_N = ${WG_TILE_N}; -const uint WG_TILE_K = ${WG_TILE_K}; - -const uint SG_GRID_X = ${SG_GRID_X}; -const uint SG_GRID_Y = ${SG_GRID_Y}; -const uint SUBGROUP_SIZE = ${SUBGROUP_SIZE}; -const uint NUM_SUBGROUPS = SG_GRID_X * SG_GRID_Y; -const uint WG_SIZE = NUM_SUBGROUPS * SUBGROUP_SIZE; - -const uint SG_TILE_M = WG_TILE_M / SG_GRID_Y; -const uint SG_TILE_N = WG_TILE_N / SG_GRID_X; -const uint MMAS_PER_SG_M = SG_TILE_M / MMA_M; -const uint MMAS_PER_SG_N = SG_TILE_N / MMA_N; - -const uint A_SLAB_INT8 = WG_TILE_M * MMA_K; -const uint B_USEFUL_U32 = MMA_K / 4u; -// Intervention A: the bank-conflict skew is what makes B's LDS strides -// non-power-of-two (B_USEFUL_U32+1 = 5, and B_SLAB_U32 = WG_TILE_N*5), forcing -// integer multiplies in every B address. B_SKEW is templated so skew=4 (round -// up to a power of two) and skew=0 (no skew, like the reference -tr's -// ROW_PAD_SH=0) can both be measured against the baseline skew of 1. -const uint B_STRIDE_U32 = B_USEFUL_U32 + ${B_SKEW}u; -const uint B_SLAB_U32 = WG_TILE_N * B_STRIDE_U32; -const uint NUM_K_SLABS = WG_TILE_K / MMA_K; - -const uint A_SLAB_U32 = A_SLAB_INT8 / 4u; -const uint A_STRIDE_U32 = MMA_K / 4u; - -// One ping-pong slice covers all K-slabs of one chunk. -const uint ASH_SLICE_U32 = NUM_K_SLABS * A_SLAB_U32; -const uint BSH_SLICE_U32 = NUM_K_SLABS * B_SLAB_U32; - -// Double-buffered MMA operand staging. -shared uint Ash_int8[2u * ASH_SLICE_U32]; -shared uint Bsh_int8[2u * BSH_SLICE_U32]; - -// Per-WG-tile-row activation params (loaded ONCE at WG start; constant -// across groups). -shared int izp_sh[WG_TILE_M]; // int32 (cast from int8 source) for broadcast -shared float ifs_sh[WG_TILE_M]; // float32 (cast from fp16 source) for broadcast - -// Per-(group, output-channel) weight params, ping-ponged by group parity. -// (For per-channel INT8 only slice 0 is ever used.) -shared float wsc_sh[2u * WG_TILE_N]; -// SUM_g wsc[g][n]*wsum[g][n] per output channel -- weight-side only, so it is -// accumulated once in the prologue. Replaces dbuf4's ping-ponged wsum_sh -// (which was 2*WG_TILE_N ints), so this is a net LDS saving. -shared float wcorr_sh[WG_TILE_N]; - -#ifdef HAS_BIAS -shared float bias_sh[WG_TILE_N]; -#endif - -#ifdef IO_TEXTURE -// Result staging for the imageStore epilogue, mirroring the fp16 kernel: -// SG_GRID_Y bands of MMA_M rows, each WG_TILE_N wide, row-major. A full -// WG_TILE_M x WG_TILE_N buffer would cost SG_GRID_Y/MMAS_PER_SG_M x more LDS -// and wreck occupancy. float16_t-typed because coopMatStore needs it. -const uint CSH_ROWS = SG_GRID_Y * MMA_M; -shared float16_t Csh_out[CSH_ROWS * WG_TILE_N]; -#endif - -// Running fp32 accumulator (across all groups). -coopmat - result[MMAS_PER_SG_M][MMAS_PER_SG_N]; - -// Per-group int32 MMA accumulator. -coopmat - accum_int32[MMAS_PER_SG_M][MMAS_PER_SG_N]; - - -// Byte-parallel int4 -> int8 widening. -// -// The four nibbles this shader needs from one packed uint are ALREADY one per -// byte (bits 3:0 of each byte for parity 0, bits 7:4 for parity 1), so all four -// can be widened at once instead of with a per-nibble -// shift/mask/bias-subtract/mask chain. -// -// For v in [0,15] the biased value is v-8. `v ^ 8` is exactly the 4-bit -// two's-complement pattern of v-8, because -8 == +8 (mod 16): -// v=0 -> 0x8 -> -8 v=7 -> 0xF -> -1 -// v=8 -> 0x0 -> 0 v=15 -> 0x7 -> +7 -// so the only remaining work is sign-extending bit 3 into bits 7:4 per byte. -// `sgn * 0x1E` does that with no cross-byte carry: 0x08 * 0x1E == 0xF0 exactly, -// and sgn is at most 0x08080808 so the product is at most 0xF0F0F0F0. -// -// A naive `nib - 0x08080808` would NOT work -- it borrows across byte lanes -// whenever a nibble is < 8. Shifts must be on uint, not int, so the >> is -// logical rather than arithmetic. -// -// ~5 ops per 4 weights vs ~22 for the per-nibble chain; bit-identical output. -uint widen_nibbles(const uint w, const uint parity) { - const uint nib = (parity == 0u) ? (w & 0x0F0F0F0Fu) : ((w >> 4u) & 0x0F0F0F0Fu); - const uint p = nib ^ 0x08080808u; - const uint sgn = p & 0x08080808u; - return p | (sgn * 0x1Eu); -} - -void main() { - const uvec2 tileID = uvec2(gl_WorkGroupID.xy); - const uvec2 warpInTile = uvec2( - gl_SubgroupID % SG_GRID_X, - gl_SubgroupID / SG_GRID_X); - - const uint K = uint(input_sizes.x); - const uint N = uint(output_sizes.x); - const uint N4 = (N + 3u) / 4u; - const uint nblocks_x_A = (K + 3u) >> 2u; - -#ifdef WEIGHT_INT4 - const uint num_groups = uint(num_groups_arg); - const uint CHUNKS_PER_GROUP = uint(K4_per_group) * 4u / WG_TILE_K; -#else - // Per-channel: a single quant "group" spanning all of K. The nested - // groups x chunks loop below collapses to a flat chunk loop, the wsum/wsc - // ping-pong never crosses a boundary, and the epilog runs exactly once. - const uint num_groups = 1u; - const uint CHUNKS_PER_GROUP = uint(num_groups_arg); -#endif - const uint num_chunks = num_groups * CHUNKS_PER_GROUP; - - const uint tile_m_start = WG_TILE_M * tileID.y; - const uint tile_n_start = WG_TILE_N * tileID.x; - - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - result[i][j] = coopmat(0.0); - accum_int32[i][j] = coopmat(0); - } - } - - // --- A staging thread map: one (m4, k4) ivec4 block per active thread --- - const uint K_BLOCKS_PER_CHUNK = WG_TILE_K >> 2u; - const uint A_ACTIVE_THREADS = (WG_TILE_M >> 2u) * K_BLOCKS_PER_CHUNK; - const uint a_m_block = gl_LocalInvocationID.x / K_BLOCKS_PER_CHUNK; - const uint a_k_block = gl_LocalInvocationID.x % K_BLOCKS_PER_CHUNK; - const bool a_active = gl_LocalInvocationID.x < A_ACTIVE_THREADS; - -#ifdef WEIGHT_INT4 - // --- B staging thread map: (block, col) slots; each slot extracts one - // ColumnMajor LDS uint (4 K-contiguous sign-extended int8) --- - const uint B_TOTAL_SLOTS = K_BLOCKS_PER_CHUNK * WG_TILE_N; - const uint B_SLOTS_PER_THREAD = B_TOTAL_SLOTS / WG_SIZE; - const uint N8_PER_TILE = WG_TILE_N >> 3u; -#else - // --- B staging thread map: one (k4, n4) ivec4 block per active thread --- - const uint B_FETCH_SLOTS = K_BLOCKS_PER_CHUNK * (WG_TILE_N >> 2u); - const uint N4_PER_TILE = WG_TILE_N >> 2u; - const uint b_k4_in_chunk = gl_LocalInvocationID.x / N4_PER_TILE; - const uint b_n_uint_col = gl_LocalInvocationID.x % N4_PER_TILE; - const bool b_active = gl_LocalInvocationID.x < B_FETCH_SLOTS; -#endif - - // Prefetch temp registers. - ivec4 temp_A; -#ifdef WEIGHT_INT4 - ivec4 temp_B[B_SLOTS_PER_THREAD]; - float temp_wsc; -#else - ivec4 temp_B; -#endif - - // ========================================================= - // PROLOGUE - // ========================================================= - if (gl_LocalInvocationID.x < (WG_TILE_M >> 2u)) { - const uint m4 = (tile_m_start >> 2u) + gl_LocalInvocationID.x; - const vec4 sc = vec4(texelFetch(t_int8_input_scales, ivec3(m4, 0, 0), 0)); - const ivec4 zp = texelFetch(t_int8_input_zps, ivec3(m4, 0, 0), 0); - const uint base = gl_LocalInvocationID.x * 4u; - ifs_sh[base + 0u] = sc.x; ifs_sh[base + 1u] = sc.y; - ifs_sh[base + 2u] = sc.z; ifs_sh[base + 3u] = sc.w; - izp_sh[base + 0u] = zp.x; izp_sh[base + 1u] = zp.y; - izp_sh[base + 2u] = zp.z; izp_sh[base + 3u] = zp.w; - } - // Group 0 weight scales -> slice 0, and the hoisted weight-side correction - // SUM_g wsc[g][n]*wsum[g][n] accumulated across ALL groups. The loop is - // prologue-only (the prologue is ~1.4% of the dynamic instruction stream), - // and it replaces per-group wsum work inside the loop body. - if (gl_LocalInvocationID.x < WG_TILE_N) { - const uint n_idx = tile_n_start + gl_LocalInvocationID.x; - f16vec4 sv0 = t_weight_scales[n_idx >> 2u]; - wsc_sh[gl_LocalInvocationID.x] = float(sv0[n_idx & 3u]); - - float corr = 0.0; - for (uint g = 0; g < num_groups; ++g) { - f16vec4 sv = t_weight_scales[g * N4 + (n_idx >> 2u)]; - corr += float(sv[n_idx & 3u]) * float(t_weight_sums[g * N + n_idx]); - } - wcorr_sh[gl_LocalInvocationID.x] = corr; - } - memoryBarrierShared(); - barrier(); - - // NOTE: dbuf4 builds izp_bcast/ifs_bcast here and keeps them live across the - // whole group loop. This variant needs them only AFTER the loop, so they are - // loaded there instead -- that is the register-pressure saving. - - // dbuf4: prefetch chunk 0 into temp registers, THEN store to slice 0 (no - // barrier here -- the main loop's first iteration barriers before - // reading slice 0). - if (a_active) { - const uint m4_global = (tile_m_start >> 2u) + a_m_block; - temp_A = t_packed_int8_input[m4_global * nblocks_x_A + a_k_block]; - } -#ifdef WEIGHT_INT4 - [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { - const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; - const uint block_in_chunk = slot >> 3u; - const uint k4_blk = block_in_chunk / N8_PER_TILE; - const uint n8_blk = (tile_n_start >> 3u) + (block_in_chunk % N8_PER_TILE); -#ifdef WEIGHT_BUFFER - temp_B[si] = t_packed_weight[(n8_blk * nblocks_x_A) + k4_blk]; -#else - temp_B[si] = texelFetch(t_packed_weight, ivec2(k4_blk, n8_blk), 0); -#endif - } -#else - if (b_active) { - const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; -#ifdef WEIGHT_BUFFER - temp_B = t_packed_weight[(b_k4_in_chunk * N4) + block_x_w]; -#else - temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, b_k4_in_chunk), 0); -#endif - } -#endif - { - // store chunk 0 -> slice 0 - if (a_active) { - const uint slab_idx = a_k_block / (MMA_K >> 2u); - const uint k_uint_in_slab = a_k_block % (MMA_K >> 2u); - const uint base_row = a_m_block * 4u; - [[unroll]] for (uint m4i = 0; m4i < 4u; ++m4i) { - Ash_int8[slab_idx * A_SLAB_U32 + (base_row + m4i) * A_STRIDE_U32 + k_uint_in_slab] = - uint(temp_A[m4i]); - } - } -#ifdef WEIGHT_INT4 - [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { - const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; - const uint block_in_chunk = slot >> 3u; - const uint col_in_block = slot & 7u; - const uint k4_in_chunk = block_in_chunk / N8_PER_TILE; - const uint n8_in_tile = block_in_chunk % N8_PER_TILE; - const uint r = col_in_block & 3u; - const uint parity = col_in_block >> 2u; - const uint n_col = n8_in_tile * 8u + r + parity * 4u; - const uint slab_idx = k4_in_chunk / (MMA_K >> 2u); - const uint k4_in_slab = k4_in_chunk % (MMA_K >> 2u); - Bsh_int8[slab_idx * B_SLAB_U32 + n_col * B_STRIDE_U32 + k4_in_slab] = - widen_nibbles(uint(temp_B[si][r]), parity); - } -#else - if (b_active) { - const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); - const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); - const uint n_col_base = b_n_uint_col * 4u; - [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { - Bsh_int8[slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab] = - uint(temp_B[n_in_blk]); - } - } -#endif - } - - // ========================================================= - // MAIN LOOP (dbuf4) — nested groups x chunks (kept nested; flattening it - // with a conditional coopmat epilog crashes the Xclipse PAL compiler at - // large spec-resolved trip counts). One barrier per chunk. Chunk - // iteration (global index `chunk`): - // 1. barrier — A/B slice (chunk%2) fully written; on the first chunk - // of group g, wsum/wsc slice (g%2) is too. - // 2. prefetch — chunk+1 (A blocks, B blocks) into temp; when chunk+1 - // starts a new group, also its wsum/wsc element. Skipped - // entirely on the final chunk. - // 3. int8 MMA — on slice (chunk%2) into accum_int32. - // 4. store — temp -> A/B slice ((chunk+1)%2), unpacking the weight; - // on a group boundary, wsum/wsc -> slice ((g+1)%2). - // The group epilog runs unconditionally at the tail of each group. - // ========================================================= - uint chunk = 0; - for (uint group_i = 0; group_i < num_groups; ++group_i) { - for (uint inner = 0; inner < CHUNKS_PER_GROUP; ++inner, ++chunk) { - const bool has_next = chunk + 1u < num_chunks; - const bool group_crossing = has_next && (inner + 1u == CHUNKS_PER_GROUP); - const uint cur_a = (chunk % 2u) * ASH_SLICE_U32; - const uint cur_b = (chunk % 2u) * BSH_SLICE_U32; - const uint nxt_a = ((chunk + 1u) % 2u) * ASH_SLICE_U32; - const uint nxt_b = ((chunk + 1u) % 2u) * BSH_SLICE_U32; - - barrier(); - - // --- 2. prefetch chunk+1 -> temp --- - if (has_next) { - const uint chunkK_nxt = (chunk + 1u) * WG_TILE_K; - if (a_active) { - const uint m4_global = (tile_m_start >> 2u) + a_m_block; - const uint k4_global = (chunkK_nxt >> 2u) + a_k_block; - temp_A = t_packed_int8_input[m4_global * nblocks_x_A + k4_global]; - } -#ifdef WEIGHT_INT4 - [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { - const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; - const uint block_in_chunk = slot >> 3u; - const uint k4_blk = (chunkK_nxt >> 2u) + block_in_chunk / N8_PER_TILE; - const uint n8_blk = (tile_n_start >> 3u) + (block_in_chunk % N8_PER_TILE); -#ifdef WEIGHT_BUFFER - temp_B[si] = t_packed_weight[(n8_blk * nblocks_x_A) + k4_blk]; -#else - temp_B[si] = texelFetch(t_packed_weight, ivec2(k4_blk, n8_blk), 0); -#endif - } - if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { - const uint n_idx = tile_n_start + gl_LocalInvocationID.x; - f16vec4 sv = t_weight_scales[(group_i + 1u) * N4 + (n_idx >> 2u)]; - temp_wsc = float(sv[n_idx & 3u]); - } -#else - if (b_active) { - const uint block_y_w = (chunkK_nxt >> 2u) + b_k4_in_chunk; - const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; -#ifdef WEIGHT_BUFFER - temp_B = t_packed_weight[(block_y_w * N4) + block_x_w]; -#else - temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, block_y_w), 0); -#endif - } -#endif - } - - // --- 3. int8 MMA on the cur slice --- - [[unroll]] for (uint k = 0; k < NUM_K_SLABS; ++k) { - const uint slab_a_base_u32 = cur_a + k * A_SLAB_U32; - const uint slab_b_base_u32 = cur_b + k * B_SLAB_U32; - - coopmat matA[MMAS_PER_SG_M]; - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - const uint row_a = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); - coopMatLoad( - matA[i], Ash_int8, - slab_a_base_u32 + row_a * A_STRIDE_U32, - A_STRIDE_U32, - gl_CooperativeMatrixLayoutRowMajor); - } - - coopmat matB; - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint col_b = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopMatLoad( - matB, Bsh_int8, - slab_b_base_u32 + col_b * B_STRIDE_U32, - B_STRIDE_U32, - gl_CooperativeMatrixLayoutColumnMajor); - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - accum_int32[i][j] = coopMatMulAdd(matA[i], matB, accum_int32[i][j]); - } - } - } - - // --- 4. store temp (chunk+1) -> nxt slice --- - if (has_next) { - if (a_active) { - const uint slab_idx = a_k_block / (MMA_K >> 2u); - const uint k_uint_in_slab = a_k_block % (MMA_K >> 2u); - const uint base_row = a_m_block * 4u; - [[unroll]] for (uint m4i = 0; m4i < 4u; ++m4i) { - Ash_int8[nxt_a + slab_idx * A_SLAB_U32 + (base_row + m4i) * A_STRIDE_U32 + k_uint_in_slab] = - uint(temp_A[m4i]); - } - } -#ifdef WEIGHT_INT4 - [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { - const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; - const uint block_in_chunk = slot >> 3u; - const uint col_in_block = slot & 7u; - const uint k4_in_chunk = block_in_chunk / N8_PER_TILE; - const uint n8_in_tile = block_in_chunk % N8_PER_TILE; - const uint r = col_in_block & 3u; - const uint parity = col_in_block >> 2u; - const uint n_col = n8_in_tile * 8u + r + parity * 4u; - const uint slab_idx = k4_in_chunk / (MMA_K >> 2u); - const uint k4_in_slab = k4_in_chunk % (MMA_K >> 2u); - Bsh_int8[nxt_b + slab_idx * B_SLAB_U32 + n_col * B_STRIDE_U32 + k4_in_slab] = - widen_nibbles(uint(temp_B[si][r]), parity); - } - if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { - const uint wbase_nxt = ((group_i + 1u) % 2u) * WG_TILE_N; - wsc_sh[wbase_nxt + gl_LocalInvocationID.x] = temp_wsc; - } -#else - if (b_active) { - const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); - const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); - const uint n_col_base = b_n_uint_col * 4u; - [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { - Bsh_int8[nxt_b + slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab] = - uint(temp_B[n_in_blk]); - } - } -#endif - } - } // chunks - - // --- Group epilog: scale-only accumulate, reset accum --- - // Just result += float(acc) * wsc. The zero-point subtract and the ifs - // multiply are hoisted out of the group loop (applied once below). - { - const uint wbase = (group_i % 2u) * WG_TILE_N; - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - - coopmat wsc_bcast; - coopMatLoad( - wsc_bcast, wsc_sh, - wbase + local_n_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutRowMajor); - - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - result[i][j] += - coopmat( - accum_int32[i][j]) * wsc_bcast; - accum_int32[i][j] = coopmat(0); - } - } - } - } // groups - - // --- Hoisted correction, applied ONCE: --------------------------------- - // result = ifs * ( result - izp * SUM_g wsc_g*wsum_g ) - // izp/ifs are loaded here rather than before the group loop so they are not - // live across it. - { - coopmat - izpf_bcast[MMAS_PER_SG_M]; - coopmat - ifs_bcast[MMAS_PER_SG_M]; - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - const uint local_m_base = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); - coopmat izp_i; - coopMatLoad( - izp_i, izp_sh, local_m_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutColumnMajor); - izpf_bcast[i] = - coopmat(izp_i); - coopMatLoad( - ifs_bcast[i], ifs_sh, local_m_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutColumnMajor); - } - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopmat wcorr_bcast; - coopMatLoad( - wcorr_bcast, wcorr_sh, local_n_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutRowMajor); - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - result[i][j] = ifs_bcast[i] * (result[i][j] - izpf_bcast[i] * wcorr_bcast); - } - } - } - - // --- Bias (optional) --- -#ifdef HAS_BIAS - if (apply_bias > 0) { - for (uint t = gl_LocalInvocationID.x; t < WG_TILE_N; t += WG_SIZE) { - bias_sh[t] = float(t_bias[tile_n_start + t]); - } - memoryBarrierShared(); - barrier(); - } -#endif - - // --- Store result tile --- - // N for the store address math MUST come from the spec constant, not the - // sizes UBO (see out_N_arg above). -#ifdef IO_TEXTURE - // Epilogue iteration i drains accumulator row-block i from EVERY subgroup - // into Csh_out at once, so the SG_GRID_Y bands it holds are disjoint global - // row ranges; the whole workgroup then imageStores them. lr / MMA_M is the - // writing subgroup's warpInTile.y, so the global row reproduces the buffer - // path's gi exactly. - // - // PORTABILITY NOTE: the barrier() in the loop body keeps this loop rolled - // despite [[unroll]], so result[i][j] IS dynamically indexed. Coopmat arrays - // are opaque per-lane storage and dynamic indexing is exactly the construct - // the Xclipse/AMD-PAL compiler has broken before -- check this first if the - // texture variants miscompile on M51. - const uint CSH_TEXELS_PER_ROW = WG_TILE_N / 4u; - const uint CSH_TEXELS = CSH_ROWS * CSH_TEXELS_PER_ROW; - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - // Guards Csh_out against the previous iteration's readers. Inert on i == 0 - // but must stay unconditional to remain workgroup-uniform. - barrier(); - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { -#ifdef HAS_BIAS - if (apply_bias > 0) { - const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopmat bias_tile; - coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); - result[i][j] += bias_tile; - } -#endif - coopmat out_tile = - coopmat(result[i][j]); - coopMatStore( - out_tile, Csh_out, - warpInTile.y * MMA_M * WG_TILE_N + - MMA_N * (MMAS_PER_SG_N * warpInTile.x + j), - WG_TILE_N, - gl_CooperativeMatrixLayoutRowMajor); - } - memoryBarrierShared(); - barrier(); - - for (uint t = gl_LocalInvocationID.x; t < CSH_TEXELS; t += WG_SIZE) { - const uint lr = t / CSH_TEXELS_PER_ROW; - const uint lc4 = t % CSH_TEXELS_PER_ROW; - const uint m = - tile_m_start + (lr / MMA_M) * SG_TILE_M + i * MMA_M + (lr % MMA_M); - const uint base = lr * WG_TILE_N + lc4 * 4u; - imageStore( - t_output, - ivec3(tile_n_start / 4u + lc4, m, 0), - vec4( - float(Csh_out[base]), - float(Csh_out[base + 1u]), - float(Csh_out[base + 2u]), - float(Csh_out[base + 3u]))); - } - } -#else - const uint N_out = uint(out_N_arg); - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint gi = tile_m_start + MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); - const uint gj = tile_n_start + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - -#ifdef HAS_BIAS - if (apply_bias > 0) { - const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopmat bias_tile; - coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); - result[i][j] += bias_tile; - } -#endif - - coopmat out_tile = - coopmat(result[i][j]); - coopMatStore( - out_tile, t_output, - gi * N_out + gj, N_out, - gl_CooperativeMatrixLayoutRowMajor); - } - } -#endif // IO_TEXTURE -} diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpb.yaml b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpb.yaml deleted file mode 100644 index b2214acce32..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpb.yaml +++ /dev/null @@ -1,42 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -# Intervention A: B LDS skew swept over {1 (baseline), 4 (power-of-two stride), -# 0 (no skew)} on top of zp-hoist + byte-parallel nibble widening, at the -# winning tile. Measured target is only 12 v_mul_lo_u32 (1.63% of the loop -# body) and it trades against an LDS category at 8.8%, so a regression is a -# plausible outcome. - -linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpb: - parameter_names_with_default_values: - PRECISION: highp - HAS_BIAS: false - IO_STORAGE: buffer - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - MMA_M: 16 - MMA_N: 16 - MMA_K: 16 - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - B_SKEW: 1 - shader_variants: - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpb_t128x64k32g42s32sk4_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - B_SKEW: 4 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpb_t128x64k32g42s32sk0_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - B_SKEW: 0 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpb_t128x64k32g42s32sk1_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - B_SKEW: 1 diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpn.glsl b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpg.glsl similarity index 85% rename from backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpn.glsl rename to backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpg.glsl index 42a813df77d..449d1710722 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpn.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpg.glsl @@ -36,6 +36,8 @@ * No new binding and no export-format change: the weight-side sum is derived * in the prologue from t_weight_scales and t_weight_sums, both already bound. * + * INTERVENTION F variant: all loop-invariant staging index arithmetic is + * hoisted out of the group loop (see "INTERVENTION F" below). * Additionally widens int4 -> int8 byte-parallel (see widen_nibbles below), * replacing the per-nibble shift/mask/bias-subtract chain. Bit-identical. * @@ -105,6 +107,15 @@ $if WEIGHT_NBITS == 4: #define WEIGHT_INT4 +// INTERVENTION G: when the A staging thread map exactly covers the workgroup +// (A_ACTIVE_THREADS == WG_SIZE) the `a_active` guard is statically always true, +// but the driver compiler does not fold it -- gl_LocalInvocationID.x's bound +// comes from a spec constant, so the comparison survives into the hot loop as a +// real branch. Set A_MAP_FULL only for tiles where the equality has been +// checked arithmetically; the yaml records the arithmetic per variant. +$if A_MAP_FULL: + #define A_ALWAYS_ACTIVE + $if HAS_BIAS: #define HAS_BIAS @@ -169,7 +180,11 @@ const uint MMAS_PER_SG_N = SG_TILE_N / MMA_N; const uint A_SLAB_INT8 = WG_TILE_M * MMA_K; const uint B_USEFUL_U32 = MMA_K / 4u; -const uint B_STRIDE_U32 = B_USEFUL_U32 + 1u; // +1 skew +// No skew + coalesced write, ported from the production dbuf4.glsl fix +// (dq8ca-dequant-unpack-ablation Addenda 6/9): stride=4 beats stride=5 on +// this hardware, and writing via the inverted contiguous-index mapping +// (see the F-hoist block below) beats the natural strided-scatter mapping. +const uint B_STRIDE_U32 = B_USEFUL_U32; const uint B_SLAB_U32 = WG_TILE_N * B_STRIDE_U32; const uint NUM_K_SLABS = WG_TILE_K / MMA_K; @@ -284,7 +299,13 @@ void main() { const uint A_ACTIVE_THREADS = (WG_TILE_M >> 2u) * K_BLOCKS_PER_CHUNK; const uint a_m_block = gl_LocalInvocationID.x / K_BLOCKS_PER_CHUNK; const uint a_k_block = gl_LocalInvocationID.x % K_BLOCKS_PER_CHUNK; +#ifdef A_ALWAYS_ACTIVE + // A_ACTIVE_THREADS == WG_SIZE for this variant's tile, so every thread stages + // A and the guard is unconditionally true. + const bool a_active = true; +#else const bool a_active = gl_LocalInvocationID.x < A_ACTIVE_THREADS; +#endif #ifdef WEIGHT_INT4 // --- B staging thread map: (block, col) slots; each slot extracts one @@ -301,6 +322,53 @@ void main() { const bool b_active = gl_LocalInvocationID.x < B_FETCH_SLOTS; #endif + // ===== INTERVENTION F: hoist loop-invariant staging index math ===== + // Every term below depends only on gl_LocalInvocationID.x, si, tile_n_start + // and tile_m_start -- none of them on `chunk` or `group_i`. In the baseline + // this same 10-line decomposition is recomputed at THREE sites (prologue + // store, in-loop prefetch, in-loop store) on every chunk iteration. The + // ablation ladder attributes -16.8% to the B staging block and shows that + // most of it is index arithmetic, not memory traffic: deleting B staging + // removed 45 of the loop's 96 address instructions and 56 of its 64 + // compare/select instructions, but only 24 memory ops. + // + // Only the k4 texel row varies with the chunk, and it varies by a constant + // stride, so the loop carries an increment instead of a recomputation. +#ifdef WEIGHT_INT4 + uint b_lds_off[B_SLOTS_PER_THREAD]; // LDS store offset within a slice + uint b_comp[B_SLOTS_PER_THREAD]; // which ivec4 component feeds this slot + uint b_par[B_SLOTS_PER_THREAD]; // nibble parity for this slot + uint b_n8blk[B_SLOTS_PER_THREAD]; // global texel column (N/8 blocks) + uint b_k4off[B_SLOTS_PER_THREAD]; // k4 offset of this slot within a chunk + // Coalesced-write inversion (ported from production dbuf4.glsl's + // bcoal_index, dq8ca-dequant-unpack-ablation Addendum 9): start from the + // CONTIGUOUS per-thread LDS index `a` this thread will write to, and derive + // which global-fetch element it needs -- instead of deriving the LDS + // address from the natural fetch grouping (which is what produced the + // B_STRIDE_U32-strided, non-coalesced writes this fix replaces). + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { + const uint a = gl_LocalInvocationID.x + si * WG_SIZE; + const uint slab_idx = a / B_SLAB_U32; + const uint local_a = a % B_SLAB_U32; + const uint n_col = local_a / B_STRIDE_U32; + const uint k4_in_slab = local_a % B_STRIDE_U32; + const uint k4_in_chunk = slab_idx * (MMA_K >> 2u) + k4_in_slab; + const uint n8_in_tile = n_col >> 3u; + const uint rem = n_col & 7u; + b_lds_off[si] = a; + b_comp[si] = rem & 3u; + b_par[si] = rem >> 2u; + b_n8blk[si] = (tile_n_start >> 3u) + n8_in_tile; + b_k4off[si] = k4_in_chunk; + } +#endif + // A staging: same argument. base_row/slab/k_uint are all invariant. + const uint a_lds_off0 = + (a_k_block / (MMA_K >> 2u)) * A_SLAB_U32 + + (a_m_block * 4u) * A_STRIDE_U32 + + (a_k_block % (MMA_K >> 2u)); + const uint a_glb_row = ((tile_m_start >> 2u) + a_m_block) * nblocks_x_A; + // Prefetch temp registers. ivec4 temp_A; #ifdef WEIGHT_INT4 @@ -355,14 +423,10 @@ void main() { } #ifdef WEIGHT_INT4 [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { - const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; - const uint block_in_chunk = slot >> 3u; - const uint k4_blk = block_in_chunk / N8_PER_TILE; - const uint n8_blk = (tile_n_start >> 3u) + (block_in_chunk % N8_PER_TILE); #ifdef WEIGHT_BUFFER - temp_B[si] = t_packed_weight[(n8_blk * nblocks_x_A) + k4_blk]; + temp_B[si] = t_packed_weight[(b_n8blk[si] * nblocks_x_A) + b_k4off[si]]; #else - temp_B[si] = texelFetch(t_packed_weight, ivec2(k4_blk, n8_blk), 0); + temp_B[si] = texelFetch(t_packed_weight, ivec2(b_k4off[si], b_n8blk[si]), 0); #endif } #else @@ -377,29 +441,18 @@ void main() { #endif { // store chunk 0 -> slice 0 - if (a_active) { - const uint slab_idx = a_k_block / (MMA_K >> 2u); - const uint k_uint_in_slab = a_k_block % (MMA_K >> 2u); - const uint base_row = a_m_block * 4u; +#ifndef A_ALWAYS_ACTIVE + if (a_active) +#endif + { [[unroll]] for (uint m4i = 0; m4i < 4u; ++m4i) { - Ash_int8[slab_idx * A_SLAB_U32 + (base_row + m4i) * A_STRIDE_U32 + k_uint_in_slab] = - uint(temp_A[m4i]); + Ash_int8[a_lds_off0 + m4i * A_STRIDE_U32] = uint(temp_A[m4i]); } } #ifdef WEIGHT_INT4 [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { - const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; - const uint block_in_chunk = slot >> 3u; - const uint col_in_block = slot & 7u; - const uint k4_in_chunk = block_in_chunk / N8_PER_TILE; - const uint n8_in_tile = block_in_chunk % N8_PER_TILE; - const uint r = col_in_block & 3u; - const uint parity = col_in_block >> 2u; - const uint n_col = n8_in_tile * 8u + r + parity * 4u; - const uint slab_idx = k4_in_chunk / (MMA_K >> 2u); - const uint k4_in_slab = k4_in_chunk % (MMA_K >> 2u); - Bsh_int8[slab_idx * B_SLAB_U32 + n_col * B_STRIDE_U32 + k4_in_slab] = - widen_nibbles(uint(temp_B[si][r]), parity); + Bsh_int8[b_lds_off[si]] = + widen_nibbles(uint(temp_B[si][b_comp[si]]), b_par[si]); } #else if (b_active) { @@ -444,21 +497,19 @@ void main() { // --- 2. prefetch chunk+1 -> temp --- if (has_next) { const uint chunkK_nxt = (chunk + 1u) * WG_TILE_K; - if (a_active) { - const uint m4_global = (tile_m_start >> 2u) + a_m_block; - const uint k4_global = (chunkK_nxt >> 2u) + a_k_block; - temp_A = t_packed_int8_input[m4_global * nblocks_x_A + k4_global]; +#ifndef A_ALWAYS_ACTIVE + if (a_active) +#endif + { + temp_A = t_packed_int8_input[a_glb_row + (chunkK_nxt >> 2u) + a_k_block]; } #ifdef WEIGHT_INT4 [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { - const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; - const uint block_in_chunk = slot >> 3u; - const uint k4_blk = (chunkK_nxt >> 2u) + block_in_chunk / N8_PER_TILE; - const uint n8_blk = (tile_n_start >> 3u) + (block_in_chunk % N8_PER_TILE); + const uint k4_blk = (chunkK_nxt >> 2u) + b_k4off[si]; #ifdef WEIGHT_BUFFER - temp_B[si] = t_packed_weight[(n8_blk * nblocks_x_A) + k4_blk]; + temp_B[si] = t_packed_weight[(b_n8blk[si] * nblocks_x_A) + k4_blk]; #else - temp_B[si] = texelFetch(t_packed_weight, ivec2(k4_blk, n8_blk), 0); + temp_B[si] = texelFetch(t_packed_weight, ivec2(k4_blk, b_n8blk[si]), 0); #endif } if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { @@ -510,29 +561,18 @@ void main() { // --- 4. store temp (chunk+1) -> nxt slice --- if (has_next) { - if (a_active) { - const uint slab_idx = a_k_block / (MMA_K >> 2u); - const uint k_uint_in_slab = a_k_block % (MMA_K >> 2u); - const uint base_row = a_m_block * 4u; +#ifndef A_ALWAYS_ACTIVE + if (a_active) +#endif + { [[unroll]] for (uint m4i = 0; m4i < 4u; ++m4i) { - Ash_int8[nxt_a + slab_idx * A_SLAB_U32 + (base_row + m4i) * A_STRIDE_U32 + k_uint_in_slab] = - uint(temp_A[m4i]); + Ash_int8[nxt_a + a_lds_off0 + m4i * A_STRIDE_U32] = uint(temp_A[m4i]); } } #ifdef WEIGHT_INT4 [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { - const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; - const uint block_in_chunk = slot >> 3u; - const uint col_in_block = slot & 7u; - const uint k4_in_chunk = block_in_chunk / N8_PER_TILE; - const uint n8_in_tile = block_in_chunk % N8_PER_TILE; - const uint r = col_in_block & 3u; - const uint parity = col_in_block >> 2u; - const uint n_col = n8_in_tile * 8u + r + parity * 4u; - const uint slab_idx = k4_in_chunk / (MMA_K >> 2u); - const uint k4_in_slab = k4_in_chunk % (MMA_K >> 2u); - Bsh_int8[nxt_b + slab_idx * B_SLAB_U32 + n_col * B_STRIDE_U32 + k4_in_slab] = - widen_nibbles(uint(temp_B[si][r]), parity); + Bsh_int8[nxt_b + b_lds_off[si]] = + widen_nibbles(uint(temp_B[si][b_comp[si]]), b_par[si]); } if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { const uint wbase_nxt = ((group_i + 1u) % 2u) * WG_TILE_N; diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpg.yaml b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpg.yaml new file mode 100644 index 00000000000..444c8347cde --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpg.yaml @@ -0,0 +1,80 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# "zp-hoisted" variants of the int8 dq8ca_q4gsw coopmat kernel: identical to +# tsweep_dbuf4 except the activation zero-point correction and the per-row +# activation scale are applied once after the group loop instead of once per +# quantization group. No new binding, no export-format change -- the weight-side +# correction sum is derived in the prologue from tensors already bound. +# +# Selected via +# ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4zpg_txkgs. +# NOT the shipped default until it passes repeated correctness runs. + +linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpg: + parameter_names_with_default_values: + PRECISION: highp + HAS_BIAS: false + IO_STORAGE: buffer + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + A_MAP_FULL: false + MMA_M: 16 + MMA_N: 16 + MMA_K: 16 + WG_TILE_M: 64 + WG_TILE_N: 32 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + shader_variants: + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpg_t128x128k64g81s64_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 8 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 64 + # A_ACTIVE_THREADS = (128>>2)*(64>>2) = 512 == WG_SIZE = 8*1*64 = 512 + A_MAP_FULL: true + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpg_t128x128k64g81s64_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 64 + SG_GRID_X: 8 + SG_GRID_Y: 1 + SUBGROUP_SIZE: 64 + A_MAP_FULL: true + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpg_t128x64k32g42s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + # A_ACTIVE_THREADS = (128>>2)*(32>>2) = 256 == WG_SIZE = 4*2*32 = 256 + A_MAP_FULL: true + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpg_t128x64k32g42s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + # small step up + # A_ACTIVE_THREADS = (128>>2)*(32>>2) = 256 == WG_SIZE = 4*2*32 = 256 + A_MAP_FULL: true diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpn.yaml b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpn.yaml deleted file mode 100644 index b40757f0fc0..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpn.yaml +++ /dev/null @@ -1,187 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -# "zp-hoisted" variants of the int8 dq8ca_q4gsw coopmat kernel: identical to -# tsweep_dbuf4 except the activation zero-point correction and the per-row -# activation scale are applied once after the group loop instead of once per -# quantization group. No new binding, no export-format change -- the weight-side -# correction sum is derived in the prologue from tensors already bound. -# -# Selected via -# ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4zpn_txkgs. -# NOT the shipped default until it passes repeated correctness runs. - -linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpn: - parameter_names_with_default_values: - PRECISION: highp - HAS_BIAS: false - IO_STORAGE: buffer - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - MMA_M: 16 - MMA_N: 16 - MMA_K: 16 - WG_TILE_M: 64 - WG_TILE_N: 32 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - shader_variants: - # production tile -- the A/B anchor against the 1.3 baseline - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpn_t64x32k32g12s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 32 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpn_t64x32k32g12s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 32 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpn_t64x32k32g12s64_buffer_buffer_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: buffer - WG_TILE_M: 64 - WG_TILE_N: 32 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - # NOTE: t128x128k64g42s32 was seeded here and then REMOVED. Its A-staging - # thread map cannot cover its blocks: A_ACTIVE_THREADS = (128/4)*(64/4) = - # 512 but WG_SIZE = 4*2*32 = 256, so half of A is never staged and the - # kernel is numerically wrong (measured 12/12 correctness failures on the - # dbuf4 equivalent). Use g44s32 for that tile instead -- 4*4*32 = 512 = A. - # deep K + big tile; the reference's winning geometry, A map valid at 512=512 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpn_t128x128k64g44s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 4 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpn_t128x128k64g44s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 4 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - # task 5.1's explicit ask - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpn_t128x128k32g42s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 32 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpn_t128x128k32g42s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 32 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - # same tile, more subgroups -> MMAS 2x2 instead of 4x2 (less accumulator pressure) - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpn_t128x128k32g44s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 32 - SG_GRID_X: 4 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpn_t128x128k32g44s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 32 - SG_GRID_X: 4 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - # moderate step up from the shipped 64x32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpn_t128x64k32g42s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpn_t128x64k32g42s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - # small step up - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpn_t64x64k32g22s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 2 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpn_t64x64k32g22s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 2 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - # big tile, shallow K -- isolates tile area from K depth - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpn_t128x128k16g42s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpn_t128x128k16g42s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp index 268d79e4919..a1a773cac1b 100644 --- a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp @@ -112,36 +112,16 @@ static const char* const kDq8caTsweepPrefixes[] = { "tsweep_dbuf2_t", "tsweep_dbuf3_t", "tsweep_dbuf4_t", - // "-tr": dbuf4's loop with the A-side global -> LDS staging swapped to - // coopMatLoad/coopMatStore (ported from shmem_double_buf4-tr.comp). Stays - // mutually exclusive with "tsweep_dbuf4_t" because position 12 is 't' vs - // '_', so prefix order here does not matter. - "tsweep_dbuf4tr_t", - // DIAGNOSTIC: same row-major layout as "-tr" but with scalar A staging, - // used to bisect a dbuf4tr correctness failure. Delete with the shader. - "tsweep_dbuf4trm_t", - "tsweep_dbuf4trd_t", - // zp-hoisted: zero-point + activation-scale correction applied once - // after the group loop instead of once per quantization group. - "tsweep_dbuf4zp_t", - // zp-hoisted + byte-parallel int4->int8 widening. - "tsweep_dbuf4zpn_t", - // + templated B LDS skew (intervention A). MEASURED-NEGATIVE: the - // bank-conflict skew is worth far more than the multiplies it costs. - "tsweep_dbuf4zpb_t", - // zpn + ARITHMETIC nibble-parity select instead of a compared select - // (intervention E of dq8ca-prefill-stall-reduction). MEASURED-NEGATIVE: - // +2.87%; the variable shift is serially dependent on parity where the - // ternary's two constant shifts were not. - "tsweep_dbuf4zpx_t", - // zpn + loop-invariant staging index arithmetic hoisted out of the group - // loop (intervention F of dq8ca-prefill-stall-reduction). - "tsweep_dbuf4zpi_t", // zpi + compile-time elision of the statically-true a_active guard - // (intervention G of dq8ca-prefill-stall-reduction). + // (intervention G of dq8ca-prefill-stall-reduction), combined with the + // dbuf4 default's own B_STRIDE_U32 skew removal + coalesced B-store + // rewrite. PROMOTED 2026-08-28 as the shipped default -- see + // dq8ca_coopmat_variant() below. "tsweep_dbuf4zpg_t", - // zpn + the a_active elision alone, no index hoist -- isolates G from F. - "tsweep_dbuf4zpk_t", + // (dq8ca-dequant-unpack-ablation Addendum 11 -- abl_aconst/abl_areadc/ + // abl_abconst -- were measurement-only variants deleted once each + // attribution was recorded; see openspec/changes/dq8ca-dequant-unpack- + // ablation/results/README.md.) // (dq8ca-dequant-unpack-ablation and its 2026-08-26 follow-ups on // xgpusw-debug08 -- abl_nodq/abl_nonib/abl_both/abl_nolds/abl_bconst/ // abl_bcont/abl_breadc/str4/str6/str8/bcoal -- were measurement-only @@ -151,6 +131,19 @@ static const char* const kDq8caTsweepPrefixes[] = { // see the B_STRIDE_U32 comment (the LDS skew removal) and the // BCoalIndex/bcoal_index comment (the coalesced B-store rewrite) in // linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4.glsl.) + // 2026-08-28: removed the now-dead prefixes for "-tr"/"-trm"/"-trd" + // (row-major-A coopMatLoad staging, deprioritized -- see + // dq8ca-uvec4-coopmat-redesign/results.md: the reference kernel's read is + // narrow too, so this direction didn't hold), "-zp"/"-zpn"/"-zpb"/"-zpx" + // (superseded single-intervention isolations from dq8ca-prefill-stall- + // reduction, folded into "-zpg" above), and "-zpi"/"-zpk" (isolation + // variants used only to attribute G vs F in that investigation). Their + // shader files were deleted with them -- this branch (release14-quant- + // shaders) ships only the validated default; the experimental siblings + // and isolation variants live on dq8ca-uvec4-redesign/dq8ca-arch-redesign + // instead. Leaving a dead prefix here means an env var could reference a + // shader that no longer exists and crash confusingly at shader lookup + // instead of failing the validation check cleanly. "tsweep_t", }; @@ -217,8 +210,30 @@ static const std::string& q4gsw_coopmat_variant() { static const std::string& dq8ca_coopmat_variant() { // Default (no ET_VK_DQ8CA_COOPMAT_VARIANT set): - // tsweep_dbuf4_t128x128k64g81s64 (WG_TILE 128x128x64, SG_GRID 8x1, wave64) -- - // PROMOTED 2026-08-26 from the prior default tsweep_dbuf4_t64x32k32g12s64, + // tsweep_dbuf4zpg_t128x64k32g42s32 (WG_TILE 128x64x32, SG_GRID 4x2, wave32) + // -- PROMOTED 2026-08-28 from the prior default + // tsweep_dbuf4_t128x128k64g81s64 (below). Combines every validated 8da4w win + // to date: zp-hoist (zero-point + // + activation-scale correction moved out of the per-group loop), the + // byte-parallel int4 unpack, static a_active branch elision, loop-invariant + // B-staging-index hoisting (interventions from dq8ca-prefill-stall- + // reduction), plus the B_STRIDE_U32 skew removal and coalesced B-store + // rewrite this file's prior default already shipped. Real, on-device, + // correctness-validated on the sibling `dq8ca-uvec4-redesign` branch (cut + // from this branch @ 8cde63eae4, same base commit) at **46.49-46.50% + // efficiency of int8 peak on 8B** -- vs. 36.97% for the prior default, a + // further +20.5/20.9/21.2% relative speedup on 8B/3B/1B respectively + // (openspec/changes/archive/2026-08-26-dq8ca-uvec4-coopmat-redesign/results.md). + // Real e2e prefill on the canonical `main-fafb46ae9c0d` SUMD driver + // (`xgpusw-debug08`/`00000b750f413c33`, maxpin 980/5333/934, ETDump- + // confirmed dispatching this exact shader on the scored 2048-token block, + // not the warmup pass): 1B 1652.95, 3B 715.58, 8B 343.80 tok/s (median of 5 + // reps each). Correctness validated 10/10 clean across 1B/3B/8B x + // buffer/texture3d on the sibling branch before promotion; re-verified on + // this branch's own build before shipping (see this function's own + // re-verification note, if present, or the promoting commit message). + // + // Prior default, PROMOTED 2026-08-26 from tsweep_dbuf4_t64x32k32g12s64, // together with the B_STRIDE_U32 skew removal in // linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4.glsl (see that file's comment). // Real, on-device measurement @@ -273,13 +288,13 @@ static const std::string& dq8ca_coopmat_variant() { static const std::string variant = [] { const char* env = std::getenv("ET_VK_DQ8CA_COOPMAT_VARIANT"); if (!env) { - return std::string("tsweep_dbuf4_t128x128k64g81s64"); + return std::string("tsweep_dbuf4zpg_t128x64k32g42s32"); } const std::string v(env); if (is_dq8ca_shippable_token(v)) { return v; } - return std::string("tsweep_dbuf4_t128x128k64g81s64"); + return std::string("tsweep_dbuf4zpg_t128x64k32g42s32"); }(); return variant; } From 1f3322ca22162effb920e72cf4c5715343a6fa71 Mon Sep 17 00:00:00 2001 From: Yanwen Xu Date: Fri, 28 Aug 2026 16:15:56 -0700 Subject: [PATCH 14/28] [ET-VK] SDPA coopmat tile-sweep mechanism, and shader-lab's sanctioned patches Adds env-var-driven tile-sweep support to the SDPA QK^T/attn*V coopmat shaders (ET_VK_SDPA_ATTN_COOPMAT_VARIANT / ET_VK_SDPA_OUT_COOPMAT_VARIANT in SDPA.cpp), mirroring the mechanism QuantizedLinear.cpp already uses for dq8ca/q4gsw. Default behavior is unchanged -- kSdpaAttnDefaultDims/ kSdpaOutDefaultDims reproduce the prior hardcoded tile exactly. Only the shipped default tile is compiled per shader (t128x64k32g22s64 for QK^T, t64x64k32g22s64 for attn*V). A wider sweep (~40 variants per shader) was tried but found unvalidated by its own admission -- that doesn't belong in this branch (final shipped shaders only), so it's trimmed here; a real sweep belongs on an experimental branch. Also includes shader-lab's two sanctioned additive patches: --json microbench output (test_llama_microbench.cpp) and ET_VK_EXTRA_INSTANCE_LAYERS (vk_api/Runtime.cpp), needed to enable Vulkan debug layers on an Android native binary. Verified on xgpusw-debug08/00000b750f413c33 (canonical main-fafb46ae9c0d driver): 48/48 correctness cases pass (3 consecutive runs, unchanged from before -- this harness's correctness gate doesn't cover SDPA), and a real e2e run confirms coherent output at the default SDPA tiles, 343.9 tok/s prefill on 8B (matching the already-promoted dq8ca default). --- .../sdpa_compute_attn_weights_coopmat.yaml | 18 +- .../ops/glsl/sdpa_compute_out_coopmat.yaml | 17 +- .../vulkan/runtime/graph/ops/impl/SDPA.cpp | 203 ++++++++-- backends/vulkan/runtime/vk_api/Runtime.cpp | 41 ++ .../test/custom_ops/test_llama_microbench.cpp | 367 +++++++++++++++++- 5 files changed, 597 insertions(+), 49 deletions(-) diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_coopmat.yaml b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_coopmat.yaml index 26b516df732..c40737e0587 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_coopmat.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_coopmat.yaml @@ -5,8 +5,14 @@ # LICENSE file in the root directory of this source tree. # KHR Cooperative Matrix SDPA QK^T kernel (prefill / LLM). Buffer-only, fp16. -# Variant name matches the dispatch: -# sdpa_compute_attn_weights_coopmat_buffer_buffer_half. +# +# Variant naming is tile-sweep-ready (ET_VK_SDPA_ATTN_COOPMAT_VARIANT, +# SDPA.cpp): sdpa_compute_attn_weights_coopmat_txkgs, +# dispatched as ..._t..._buffer_buffer_half. Only the shipped default is +# built here -- a wider sweep was tried and found unvalidated (dozens of +# untested tiles), which doesn't belong in this branch (final shipped +# shaders only); run a real sweep on an experimental branch before adding +# more variants here. sdpa_compute_attn_weights_coopmat: parameter_names_with_default_values: @@ -34,4 +40,10 @@ sdpa_compute_attn_weights_coopmat: DTYPE: - VALUE: half shader_variants: - - NAME: sdpa_compute_attn_weights_coopmat + - NAME: sdpa_compute_attn_weights_coopmat_t128x64k32g22s64 + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coopmat.yaml b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coopmat.yaml index e5576003be7..b7859e358cd 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coopmat.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coopmat.yaml @@ -5,7 +5,14 @@ # LICENSE file in the root directory of this source tree. # KHR Cooperative Matrix SDPA attn*V kernel (prefill / LLM). Buffer-only, fp16. -# Variant name matches the dispatch: sdpa_compute_out_coopmat_buffer_buffer_half. +# +# Variant naming is tile-sweep-ready (ET_VK_SDPA_OUT_COOPMAT_VARIANT, +# SDPA.cpp): sdpa_compute_out_coopmat_txkgs, +# dispatched as ..._t..._buffer_buffer_half. Only the shipped default is +# built here -- a wider sweep was tried and found unvalidated (dozens of +# untested tiles), which doesn't belong in this branch (final shipped +# shaders only); run a real sweep on an experimental branch before adding +# more variants here. sdpa_compute_out_coopmat: parameter_names_with_default_values: @@ -30,4 +37,10 @@ sdpa_compute_out_coopmat: DTYPE: - VALUE: half shader_variants: - - NAME: sdpa_compute_out_coopmat + - NAME: sdpa_compute_out_coopmat_t64x64k32g22s64 + WG_TILE_M: 64 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 2 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 diff --git a/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp b/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp index 29469cc0770..9fe48b6fdec 100644 --- a/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp @@ -21,8 +21,10 @@ #include #include +#include #include #include +#include namespace vkcompute { @@ -190,15 +192,112 @@ static inline SDPAMode mode_of(const std::vector& resize_args) { // capability-eligible devices; ET_VK_DISABLE_COOPMAT remains the kill switch // (shared with the q4gsw linear coopmat path). // -constexpr uint32_t kSdpaCmTileM = 64; -constexpr uint32_t kSdpaCmTileN = 64; -constexpr uint32_t kSdpaCmTileK = 32; -constexpr uint32_t kSdpaCmInvocations = 256; -// QK^T uses a 128-tall M-tile (the 128x64 tile-sweep optimum that still fits -// the masked shader's shared memory); attn*V keeps the 64x64 geometry. Same -// WG_SIZE (256 = 2x2 subgroups x 64), so only the M-tile count differs for -// QK^T. -constexpr uint32_t kSdpaCmQkTileM = 128; +// SDPA coopmat tile-sweep variants (mirrors QuantizedLinear.cpp's +// ET_VK_Q4GSW_COOPMAT_VARIANT / ET_VK_DQ8CA_COOPMAT_VARIANT mechanism). +// sdpa_compute_attn_weights_coopmat (QK^T) and sdpa_compute_out_coopmat +// (attn*V) are independent shaders with independent tile geometries, so each +// gets its own env var. Token format +// "txkgs" +// -- no dbuf-namespace prefix needed since SDPA has no loop-structure +// variants (unlike q4gsw/dq8ca's dbuf1-4), so this drops the "tsweep_" +// prefix QuantizedLinear.cpp's kTsweepPrefixes uses for that +// disambiguation. +struct SdpaCoopmatTileDims { + uint32_t m; + uint32_t n; + uint32_t k; + uint32_t sgx; + uint32_t sgy; + uint32_t sub; +}; +inline uint32_t sdpa_wg_size(const SdpaCoopmatTileDims& d) { + return d.sgx * d.sgy * d.sub; +} + +// Current shipped tiles (sdpa_compute_attn_weights_coopmat.yaml / +// sdpa_compute_out_coopmat.yaml parameter_names_with_default_values), +// expressed as tokens so the default path and an explicit env-var override +// run through the identical parse+dispatch code (QuantizedLinear.cpp's +// q4gsw crash -- see 3ceeefc269 -- was exactly a bare/no-token default +// skipping this). QK^T's 128-tall M-tile is the 128x64 tile-sweep optimum +// that still fits the masked shader's shared memory (a naive reuse of the +// generic-matmul sweep's 128x128 winner needs ~50KB LDS once the +// causal-mask Csh scratch is added -- overflows M51); attn*V keeps 64x64 +// (no Csh, smaller budget). Both use the same WG_SIZE (256 = 2x2 subgroups +// x 64), so only the M-tile differs between them. +constexpr SdpaCoopmatTileDims kSdpaAttnDefaultDims = {128, 64, 32, 2, 2, 64}; +constexpr SdpaCoopmatTileDims kSdpaOutDefaultDims = {64, 64, 32, 2, 2, 64}; + +static bool is_recognized_sdpa_coopmat_variant_token(const std::string& v) { + if (v.size() < 2 || v[0] != 't' || !std::isdigit((unsigned char)v[1])) { + return false; + } + return v.find('x') != std::string::npos && v.find('k') != std::string::npos && + v.find('g') != std::string::npos && v.find('s') != std::string::npos; +} + +// Parses "txkgs" -> dims. Returns fallback unchanged +// for an unrecognized token (mirrors QuantizedLinear.cpp's +// parse_tsweep_tile). +static SdpaCoopmatTileDims parse_sdpa_tsweep_tile( + const std::string& variant, + const SdpaCoopmatTileDims& fallback) { + if (!is_recognized_sdpa_coopmat_variant_token(variant)) { + return fallback; + } + const size_t t_pos = 1; // token always starts with 't' + const size_t x_pos = variant.find('x', t_pos); + const size_t k_pos = variant.find('k', x_pos); + const size_t g_pos = variant.find('g', k_pos); + const size_t s_pos = variant.find('s', g_pos); + const uint32_t m = std::stoul(variant.substr(t_pos, x_pos - t_pos)); + const uint32_t n = std::stoul(variant.substr(x_pos + 1, k_pos - x_pos - 1)); + const uint32_t k = std::stoul(variant.substr(k_pos + 1, g_pos - k_pos - 1)); + const std::string grid = variant.substr(g_pos + 1, s_pos - g_pos - 1); + const uint32_t sgx = grid[0] - '0'; + const uint32_t sgy = grid[1] - '0'; + const uint32_t sub = std::stoul(variant.substr(s_pos + 1)); + return {m, n, k, sgx, sgy, sub}; +} + +static const std::string& sdpa_attn_weights_coopmat_variant() { + static const std::string variant = [] { + const char* env = std::getenv("ET_VK_SDPA_ATTN_COOPMAT_VARIANT"); + if (!env) { + return std::string("t128x64k32g22s64"); + } + const std::string v(env); + if (is_recognized_sdpa_coopmat_variant_token(v)) { + return v; + } + return std::string("t128x64k32g22s64"); + }(); + return variant; +} + +static const std::string& sdpa_out_coopmat_variant() { + static const std::string variant = [] { + const char* env = std::getenv("ET_VK_SDPA_OUT_COOPMAT_VARIANT"); + if (!env) { + return std::string("t64x64k32g22s64"); + } + const std::string v(env); + if (is_recognized_sdpa_coopmat_variant_token(v)) { + return v; + } + return std::string("t64x64k32g22s64"); + }(); + return variant; +} + +static SdpaCoopmatTileDims sdpa_attn_tile_dims() { + return parse_sdpa_tsweep_tile( + sdpa_attn_weights_coopmat_variant(), kSdpaAttnDefaultDims); +} +static SdpaCoopmatTileDims sdpa_out_tile_dims() { + return parse_sdpa_tsweep_tile( + sdpa_out_coopmat_variant(), kSdpaOutDefaultDims); +} static bool sdpa_coopmat_not_disabled() { return std::getenv("ET_VK_DISABLE_COOPMAT") == nullptr; @@ -224,10 +323,14 @@ static bool sdpa_buf_half(ComputeGraph* graph, const ValueRef t) { graph->dtype_of(t) == vkapi::kHalf; } -static inline bool sdpa_cm_aligned(int64_t m, int64_t n, int64_t k) { - return m % static_cast(kSdpaCmTileM) == 0 && - n % static_cast(kSdpaCmTileN) == 0 && - k % static_cast(kSdpaCmTileK) == 0; +static inline bool sdpa_cm_aligned( + int64_t m, + int64_t n, + int64_t k, + const SdpaCoopmatTileDims& dims) { + return m % static_cast(dims.m) == 0 && + n % static_cast(dims.n) == 0 && + k % static_cast(dims.k) == 0; } static inline bool is_sdpa_coopmat(const vkapi::ShaderInfo& shader) { @@ -251,10 +354,12 @@ vkapi::ShaderInfo pick_sdpa_qk_shader( sdpa_buf_half(graph, attn_weights)) { const SDPADims d = compute_sdpa_dims( *graph, q_projected, k_cache, resize_args.at(2), SDPAMode::LLM); - if (d.S % static_cast(kSdpaCmQkTileM) == 0 && - d.context_len % static_cast(kSdpaCmTileN) == 0 && - d.D % static_cast(kSdpaCmTileK) == 0) { - std::string shader_name = "sdpa_compute_attn_weights_coopmat"; + const SdpaCoopmatTileDims attn_dims = sdpa_attn_tile_dims(); + if (d.S % static_cast(attn_dims.m) == 0 && + d.context_len % static_cast(attn_dims.n) == 0 && + d.D % static_cast(attn_dims.k) == 0) { + std::string shader_name = "sdpa_compute_attn_weights_coopmat_" + + sdpa_attn_weights_coopmat_variant(); add_storage_type_suffix( shader_name, graph->storage_type_of(q_projected)); add_storage_type_suffix(shader_name, graph->storage_type_of(k_cache)); @@ -297,15 +402,18 @@ utils::uvec3 pick_sdpa_qk_global_wg_size( const SDPADims d = compute_sdpa_dims(*graph, q, k, input_pos_symint, mode); if (is_sdpa_coopmat(shader)) { - // One workgroup per 64x64 output tile (N = context_len, M = S); - // *kSdpaCmInvocations cancels the framework div_up against local x. z - // carries the head index. + // One workgroup per output tile (N = context_len, M = S); *wg_size + // cancels the framework div_up against local x. z carries the head + // index. Tile dims come from the active ET_VK_SDPA_ATTN_COOPMAT_VARIANT + // (default: unchanged from the old kSdpaCmQkTileM/kSdpaCmTileN + // constants). + const SdpaCoopmatTileDims dims = sdpa_attn_tile_dims(); const uint32_t num_tiles_n = - utils::div_up(static_cast(d.context_len), kSdpaCmTileN); + utils::div_up(static_cast(d.context_len), dims.n); const uint32_t num_tiles_m = - utils::div_up(static_cast(d.S), kSdpaCmQkTileM); + utils::div_up(static_cast(d.S), dims.m); return { - num_tiles_n * kSdpaCmInvocations, + num_tiles_n * sdpa_wg_size(dims), num_tiles_m, static_cast(d.H * d.B)}; } @@ -325,9 +433,11 @@ utils::uvec3 pick_sdpa_qk_local_wg_size( const SDPAMode mode = mode_of(resize_args); if (mode == SDPAMode::LLM) { // _coopmat must be checked before _coop (the former contains the latter as - // a substring); the coopmat shaders use a flat 256-lane workgroup. + // a substring); the coopmat shaders use a flat workgroup sized by the + // active ET_VK_SDPA_ATTN_COOPMAT_VARIANT (default: 256 = 2x2 subgroups x + // 64, unchanged from the old kSdpaCmInvocations constant). if (is_sdpa_coopmat(shader)) { - return {kSdpaCmInvocations, 1, 1}; + return {sdpa_wg_size(sdpa_attn_tile_dims()), 1, 1}; } const bool use_coop_algorithm = shader.kernel_name.find("_coop") != std::string::npos; @@ -402,8 +512,11 @@ vkapi::ShaderInfo pick_sdpa_av_shader( resize_args.at(1), resize_args.at(2), SDPAMode::LLM); - if (sdpa_cm_aligned(/*m=*/d.S, /*n=*/d.D, /*k=*/d.context_len)) { - std::string shader_name = "sdpa_compute_out_coopmat"; + const SdpaCoopmatTileDims out_dims = sdpa_out_tile_dims(); + if (sdpa_cm_aligned( + /*m=*/d.S, /*n=*/d.D, /*k=*/d.context_len, out_dims)) { + std::string shader_name = + "sdpa_compute_out_coopmat_" + sdpa_out_coopmat_variant(); add_storage_type_suffix(shader_name, graph->storage_type_of(out)); add_storage_type_suffix(shader_name, graph->storage_type_of(v_cache)); add_dtype_suffix(shader_name, graph->dtype_of(out)); @@ -440,13 +553,16 @@ utils::uvec3 pick_sdpa_av_global_wg_size( const SDPADims d = compute_sdpa_dims(*graph, q, k, input_pos_symint, mode); if (is_sdpa_coopmat(shader)) { - // One workgroup per 64x64 output tile (N = head_dim, M = S). z = head. + // One workgroup per output tile (N = head_dim, M = S). z = head. Tile + // dims come from the active ET_VK_SDPA_OUT_COOPMAT_VARIANT (default: + // unchanged from the old kSdpaCmTileM/kSdpaCmTileN constants). + const SdpaCoopmatTileDims dims = sdpa_out_tile_dims(); const uint32_t num_tiles_n = - utils::div_up(static_cast(d.D), kSdpaCmTileN); + utils::div_up(static_cast(d.D), dims.n); const uint32_t num_tiles_m = - utils::div_up(static_cast(d.S), kSdpaCmTileM); + utils::div_up(static_cast(d.S), dims.m); return { - num_tiles_n * kSdpaCmInvocations, + num_tiles_n * sdpa_wg_size(dims), num_tiles_m, static_cast(d.H * d.B)}; } @@ -465,9 +581,11 @@ utils::uvec3 pick_sdpa_av_local_wg_size( const SDPAMode mode = mode_of(resize_args); if (mode == SDPAMode::LLM) { // _coopmat must be checked before _coop (the former contains the latter as - // a substring); the coopmat shaders use a flat 256-lane workgroup. + // a substring); the coopmat shaders use a flat workgroup sized by the + // active ET_VK_SDPA_OUT_COOPMAT_VARIANT (default: 256 = 2x2 subgroups x + // 64, unchanged from the old kSdpaCmInvocations constant). if (is_sdpa_coopmat(shader)) { - return {kSdpaCmInvocations, 1, 1}; + return {sdpa_wg_size(sdpa_out_tile_dims()), 1, 1}; } const bool use_coop_algorithm = shader.kernel_name.find("_coop") != std::string::npos; @@ -562,10 +680,13 @@ void add_sdpa_compute_attn_weights_node( {}, // Specialization Constants: {inv_scale (id 3), num_k_chunks (id 4)}. // num_k_chunks = head_dim / WG_TILE_K is static and consumed only by the - // coopmat QK^T variant; the tiled/coop variants declare only id 3 and - // ignore the trailing entry. + // coopmat QK^T variant (WG_TILE_K from the active + // ET_VK_SDPA_ATTN_COOPMAT_VARIANT); the tiled/coop variants declare + // only id 3 and ignore the trailing entry -- safe to compute + // unconditionally even when coopmat doesn't end up firing. {scale_val, - graph.size_at(-1, q) / static_cast(kSdpaCmTileK)}, + graph.size_at(-1, q) / + static_cast(sdpa_attn_tile_dims().k)}, // Resize Args: [q, k, input_pos_symint_or_dummy, mode] {q, k, input_pos_symint, mode_ref}, // Resizing Logic @@ -649,14 +770,16 @@ void add_sdpa_compute_out_node( // variant — the tiled/coop variants ignore the trailing entries, and id 3 is // the inv_scale slot the decode _coop shader reads, kept at 1.0 = no-op). // num_k_chunks uses max_context_len (the loop bound is a spec const per the - // Xclipse bug); beyond-context chunks are zero-staged in the shader. - // Values are meaningful only in LLM mode; in FUSED they are ignored. + // Xclipse bug) and the active ET_VK_SDPA_OUT_COOPMAT_VARIANT's WG_TILE_K + // (independent of QK^T's own K-tile -- these are two separately-swept + // shaders); beyond-context chunks are zero-staged in the shader. Values + // are meaningful only in LLM mode; in FUSED they are ignored. const int32_t cm_head_dim = graph.size_at(-1, q); const int32_t cm_num_q_heads = graph.size_at(-2, q); const int32_t cm_max_context = graph.size_at(-3, v); + const int32_t cm_out_tile_k = static_cast(sdpa_out_tile_dims().k); const int32_t cm_num_k_chunks = - (cm_max_context + static_cast(kSdpaCmTileK) - 1) / - static_cast(kSdpaCmTileK); + (cm_max_context + cm_out_tile_k - 1) / cm_out_tile_k; const int32_t cm_out_row_stride = cm_num_q_heads * cm_head_dim; graph.execute_nodes().emplace_back(new DynamicDispatchNode( diff --git a/backends/vulkan/runtime/vk_api/Runtime.cpp b/backends/vulkan/runtime/vk_api/Runtime.cpp index abf51744a95..ad0d38b984f 100644 --- a/backends/vulkan/runtime/vk_api/Runtime.cpp +++ b/backends/vulkan/runtime/vk_api/Runtime.cpp @@ -104,6 +104,47 @@ VkInstance create_instance(const RuntimeConfig& config) { #endif /* VK_EXT_debug_report */ } + // ET_VK_EXTRA_INSTANCE_LAYERS: comma-separated instance layer names to + // enable, in addition to whatever the config already asks for. + // + // Android's Vulkan loader does not honour VK_INSTANCE_LAYERS, + // /data/local/debug/vulkan or the debug.vulkan.layers property for a native + // executable, so a layer such as VK_LAYER_LUNARG_api_dump can only be turned + // on by the application naming it here. Without this hook there is no way to + // capture an api_dump trace of a Vulkan-backend binary on device at all. + // + // Unset -- the default -- leaves the enabled layer list byte-identical to + // before, and find_requested_layers_and_extensions() below already filters + // out any name the loader does not report, so an unavailable or misspelled + // layer is skipped rather than fatal. + // + // The strings are held in a function-local static because the pointers + // handed to VkInstanceCreateInfo must outlive this scope; they are only + // read after the vector has stopped growing. + static std::vector extra_instance_layers; + if (const char* extra_layers_env = + std::getenv("ET_VK_EXTRA_INSTANCE_LAYERS")) { + const std::string spec(extra_layers_env); + size_t start = 0; + while (start <= spec.size()) { + const size_t comma = spec.find(',', start); + const size_t end = (comma == std::string::npos) ? spec.size() : comma; + std::string name = spec.substr(start, end - start); + const size_t first = name.find_first_not_of(" \t"); + const size_t last = name.find_last_not_of(" \t"); + if (first != std::string::npos) { + extra_instance_layers.push_back(name.substr(first, last - first + 1)); + } + if (comma == std::string::npos) { + break; + } + start = comma + 1; + } + for (const std::string& name : extra_instance_layers) { + requested_layers.push_back(name.c_str()); + } + } + VkInstanceCreateFlags instance_flags = 0; #ifdef __APPLE__ instance_flags |= VK_INSTANCE_CREATE_ENUMERATE_PORTABILITY_BIT_KHR; diff --git a/backends/vulkan/test/custom_ops/test_llama_microbench.cpp b/backends/vulkan/test/custom_ops/test_llama_microbench.cpp index d566bc5eb2b..30ba08e33c1 100644 --- a/backends/vulkan/test/custom_ops/test_llama_microbench.cpp +++ b/backends/vulkan/test/custom_ops/test_llama_microbench.cpp @@ -77,6 +77,7 @@ #include #include #include +#include #include #include #include @@ -118,9 +119,18 @@ struct Record { // overhead); kernel_us isolates the linear kernel itself so the two // schemes' shader-level numbers stay comparable. -1 for sdpa rows. float kernel_us = -1.0f; + // Median and coefficient of variation of the linear kernel's own + // per-iteration timings. The text report has always shown mean +/- stdev; + // these are for --json, where a consumer ranking candidates needs a robust + // centre and an explicit noise figure rather than a mean that one outlier + // can move. -1 where the suite does not produce per-iteration samples. + float kernel_median_us = -1.0f; + float kernel_cov = -1.0f; float gflops = -1.0f; // no SDPA meaning (-1 sentinel, per specs/021) std::string dispatch = "not_applicable"; std::string correctness = "SKIPPED"; + // Failure text, for the cases where the framework threw. Empty otherwise. + std::string detail_note; bool ok = false; }; @@ -128,6 +138,14 @@ std::vector g_records; // specs/021 (research.md Decision 1): shared unified RESULT,... line -- // 12 shared fields, then suite-specific extras. +// Append to the record set WITHOUT printing. emit() below both records and +// prints a RESULT line; the correctness matrix must not gain one, because the +// text output has to stay byte-identical -- but --json still needs its +// verdicts, which are the whole point of a correctness gate. +void record_only(const Record& r) { + g_records.push_back(r); +} + void emit(const Record& r) { g_records.push_back(r); std::cout << "RESULT," << r.suite << "," << r.model << "," << r.scheme << "," @@ -312,6 +330,14 @@ std::vector as_f(const ValueSpec& s) { } return o; } +// Opt-in escape hatch for run_production_diff() below: false everywhere else, +// so the perf sweep's SKIPPED-via-throw behavior at production M/N/K is +// unchanged (et-microbench-correctness-gate-blind-above-256 still applies to +// every existing call site). --production-diff sets this for the duration of +// its own run only, accepting the O(M*N*K) CPU cost as a one-shot diagnostic, +// not something to pay on every correctness/perf invocation. +bool g_allow_large_reference = false; + void bench_reference(TestCase& tc) { const std::string op = tc.operator_name(); const bool dq8ca = op.find("dq8ca") != std::string::npos; @@ -331,7 +357,7 @@ void bench_reference(TestCase& tc) { // production-K correctness cases (K=2048/4096, M/N still <=256): without // this, those cases silently throw here and get marked SKIPPED, giving a // false impression of "validated" when no reference was ever computed. - if (M > 256 || N > 256 || K > 4096) { + if (!g_allow_large_reference && (M > 256 || N > 256 || K > 4096)) { throw std::invalid_argument("ref: too big"); } // input layouts: weight-only = {in, w, w_scales, [group], bias}; @@ -649,6 +675,70 @@ float linear_kernel_us(const BenchmarkResult& r) { return us; } +// Median and CoV of the linear kernel's per-iteration timings, from the same +// ShaderTiming entries linear_kernel_us() averages. +void linear_kernel_dist( + const BenchmarkResult& r, + float* median_us, + float* cov) { + *median_us = -1.0f; + *cov = -1.0f; + std::vector samples; + for (const auto& st : r.get_shader_timings()) { + if (st.shader_name.find("linear_") != std::string::npos) { + samples = st.iter_timings_us; + } + } + if (samples.empty()) { + return; + } + std::vector sorted = samples; + std::sort(sorted.begin(), sorted.end()); + const size_t n = sorted.size(); + *median_us = + (n % 2) ? sorted[n / 2] : 0.5f * (sorted[n / 2 - 1] + sorted[n / 2]); + double sum = 0.0; + for (float v : samples) { + sum += v; + } + const double mean = sum / n; + double var = 0.0; + for (float v : samples) { + var += (v - mean) * (v - mean); + } + var /= n; + *cov = mean > 0.0 ? static_cast(std::sqrt(var) / mean) : -1.0f; +} + +// Recover (M, K, N) from a case name of the form "..._M_K_N_...". +void parse_mkn_from_case_name( + const std::string& name, + int64_t* M, + int64_t* K, + int64_t* N) { + auto grab = [&name](char tag) -> int64_t { + const std::string needle = std::string("_") + tag; + for (size_t i = 0; i + needle.size() < name.size(); ++i) { + if (name.compare(i, needle.size(), needle) != 0) { + continue; + } + size_t j = i + needle.size(); + if (j >= name.size() || !std::isdigit((unsigned char)name[j])) { + continue; + } + int64_t v = 0; + while (j < name.size() && std::isdigit((unsigned char)name[j])) { + v = v * 10 + (name[j++] - '0'); + } + return v; + } + return 0; + }; + *M = grab('M'); + *K = grab('K'); + *N = grab('N'); +} + std::string kernel_class(const std::string& kernel) { // _coopmat must be checked before _coop (substring). if (kernel.find("_coopmat") != std::string::npos) { @@ -690,9 +780,34 @@ bool run_linear_correctness(const CaseFilter& filter) { failed_names.push_back(tc.name()); } results.push_back(res[0]); + Record rec; + rec.suite = "correctness"; + rec.op = tc.name(); + // make_linear_case() builds the name from M/K/N, so recovering them + // from it keeps the JSON's shape fields populated for correctness + // rows too -- a consumer cannot tell whether a fallback was legal + // without knowing the shape the case ran at. + parse_mkn_from_case_name(tc.name(), &rec.M, &rec.K, &rec.N); + rec.kernel = linear_kernel(res[0]); + rec.variant = kernel_class(rec.kernel); + rec.kernel_us = linear_kernel_us(res[0]); + linear_kernel_dist(res[0], &rec.kernel_median_us, &rec.kernel_cov); + const auto st = res[0].get_correctness_status(); + rec.correctness = st == CorrectnessStatus::PASSED ? "PASSED" + : st == CorrectnessStatus::FAILED ? "FAILED" + : "SKIPPED"; + rec.ok = st != CorrectnessStatus::FAILED; + record_only(rec); } } catch (const std::exception& e) { failed_names.push_back(tc.name()); + Record rec; + rec.suite = "correctness"; + rec.op = tc.name(); + rec.correctness = "FAILED"; + rec.detail_note = e.what(); + rec.ok = false; + record_only(rec); std::cout << "[correctness] " << tc.name() << " FAILED: " << e.what() << "\n"; } @@ -736,6 +851,98 @@ bool run_linear_correctness(const CaseFilter& filter) { return all_ok; } +// --production-diff: a real element-wise numeric check at the actual 8B +// prefill M=2048 GEMM shapes, which run_linear_correctness's gate never +// covers (et-microbench-correctness-gate-blind-above-256 -- bench_reference +// throws above M/N=256 and those cases are marked SKIPPED, not validated). +// Every shader change promoted to the shipped dq8ca default so far +// (dq8ca-dequant-unpack-ablation Addenda 7-9) has only ever cleared 10+ +// consecutive PASSED runs at M/N<=256; this is the check that was missing +// before trusting any of it at production size. +// +// Deliberately dq8ca (8da4w) only, buffer storage only (the shipped coopmat +// dispatch's storage mode) -- matches this investigation's current scope. +// Uses the same well-conditioned deterministic data generator as the +// existing correctness matrix (positive values, no fp16-cancellation noise) +// so a mismatch here means the shader's addressing/arithmetic is wrong at +// this shape, not that the reference itself is numerically unstable. +// +// Cost: O(M*N*K) CPU reference per shape, unthrottled via +// g_allow_large_reference -- on the order of minutes total for all four 8B +// shapes. That is acceptable for a one-shot diagnostic; it must never run as +// part of the default correctness gate or perf sweep. +bool run_production_diff() { + const LinearModel* model = nullptr; + for (const auto& m : kLinearModels) { + if (std::string(m.model) == "llama-3.1-8b") { + model = &m; + } + } + if (model == nullptr) { + std::cout << "[production-diff] llama-3.1-8b not found in kLinearModels\n"; + return false; + } + g_allow_large_reference = true; + bool all_ok = true; + for (const auto& op_shape : model->ops) { + LinearConfig cfg{ + /*M=*/2048, + op_shape.K, + op_shape.N, + /*group_size=*/g_group, + /*op_name=*/"linear_dq8ca_q4gsw", + /*batch=*/0, + /*model=*/model->model, + /*regime=*/"prefill", + /*op_label=*/op_shape.op_label}; + TestCase tc = make_deterministic_correctness_case( + cfg, "linear_dq8ca_q4gsw", utils::kBuffer); + std::cout << "[production-diff] " << tc.name() + << " (M=2048, K=" << op_shape.K << ", N=" << op_shape.N + << ", group_size=" << g_group << ")\n"; + try { + auto res = execute_test_cases( + [&tc]() { return std::vector{tc}; }, + flop_calc, + "LlamaMicrobenchProductionDiff", + kWarmupRuns, + kTimedRuns, + bench_reference); + if (res.empty()) { + std::cout << "[production-diff] " << tc.name() + << ": no result produced\n"; + all_ok = false; + continue; + } + const std::string shader_name = linear_kernel(res[0]); + const bool coopmat_fired = + shader_name.find("coopmat") != std::string::npos; + const bool passed = + res[0].get_correctness_status() == CorrectnessStatus::PASSED; + all_ok = all_ok && coopmat_fired && passed; + std::cout << "[production-diff] " << tc.name() << " -> " << shader_name + << (coopmat_fired ? " (coopmat dispatched)" + : " (NOT coopmat -- fallback, cannot " + "validate the shader under test)") + << ", correctness=" + << (passed ? "PASSED" + : (res[0].get_correctness_status() == + CorrectnessStatus::FAILED + ? "FAILED (see mismatch detail above)" + : "SKIPPED")) + << "\n"; + } catch (const std::exception& e) { + std::cout << "[production-diff] " << tc.name() << " threw: " << e.what() + << "\n"; + all_ok = false; + } + } + g_allow_large_reference = false; + std::cout << "[production-diff] " << (all_ok ? "ALL PASSED" : "FAILED") + << " (4 shapes, M=2048, 8da4w, buffer)\n"; + return all_ok; +} + struct PerfCase { LinearConfig cfg; utils::StorageType storage; @@ -820,6 +1027,7 @@ void run_linear_suite(const std::string& suite, const CaseFilter& filter) { rec.stdev_us = res[0].get_std_dev_us(); rec.kernel = linear_kernel(res[0]); rec.kernel_us = linear_kernel_us(res[0]); + linear_kernel_dist(res[0], &rec.kernel_median_us, &rec.kernel_cov); rec.variant = kernel_class(rec.kernel); rec.gflops = rec.mean_us > 0 ? (2.0f * cfg.M * cfg.N * cfg.K) / (rec.mean_us * 1e3f) @@ -1447,6 +1655,100 @@ std::string fmt_x(float x) { // Prints the raw-results table, the per-site WMMA speedups, and the // geomeans. Returns false if any expected coopmat site failed to speed up // AND failed to dispatch -- dispatch anomalies, not slowness, fail the run. +// --- additive machine-readable output (--json) ---------------------------- +// +// The text report is for a human reading a terminal; this is for the L2 stage +// of an automated ladder, which needs per-case median, CoV, GFLOP/s, the +// kernel that actually dispatched, and the correctness verdict, all keyed so a +// candidate can be matched to the variant it was selected as. +// +// Deliberately additive: it prints nothing unless --json is given, changes no +// existing line, and computes nothing the text path did not already compute. + +std::string json_escape(const std::string& in) { + std::string out; + for (char c : in) { + switch (c) { + case '"': + out += "\\\""; + break; + case '\\': + out += "\\\\"; + break; + case '\n': + out += "\\n"; + break; + case '\r': + out += "\\r"; + break; + case '\t': + out += "\\t"; + break; + default: + if (static_cast(c) < 0x20) { + char buf[8]; + snprintf(buf, sizeof(buf), "\\u%04x", c); + out += buf; + } else { + out += c; + } + } + } + return out; +} + +// JSON has no NaN or Infinity, and the -1 sentinels mean "not applicable" +// rather than "minus one"; both become null so a consumer cannot mistake +// either for a measurement. +std::string json_num(float v) { + if (std::isnan(v) || std::isinf(v) || v < 0.0f) { + return "null"; + } + std::ostringstream o; + o << v; + return o.str(); +} + +void print_json_report(std::ostream& out) { + out << "{\n \"schema\": \"test_llama_microbench.v1\""; + { + const auto* adapter = api::context()->adapter_ptr(); + out << ",\n \"device\": \"" << json_escape(adapter->device_name()) << "\"" + << ",\n \"subgroup_size\": " << adapter->subgroup_size() + << ",\n \"timestamp_period_ns\": " << adapter->timestamp_period() + << ",\n \"cooperative_matrix\": " + << (adapter->supports_cooperative_matrix() ? "true" : "false"); + } + out << ",\n \"warmup_runs\": " << kWarmupRuns + << ",\n \"timed_runs\": " << kTimedRuns + << ",\n \"group_size\": " << g_group << ",\n \"cases\": [\n"; + for (size_t i = 0; i < g_records.size(); ++i) { + const Record& r = g_records[i]; + out << " {" << "\"suite\": \"" << json_escape(r.suite) << "\"" + << ", \"model\": \"" << json_escape(r.model) << "\"" + << ", \"scheme\": \"" << json_escape(r.scheme) << "\"" + << ", \"regime\": \"" << json_escape(r.regime) << "\"" << ", \"op\": \"" + << json_escape(r.op) << "\"" << ", \"storage\": \"" + << json_escape(r.storage) << "\"" << ", \"variant\": \"" + << json_escape(r.variant) << "\"" << ", \"kernel\": \"" + << json_escape(r.kernel) << "\"" << ", \"M\": " << r.M + << ", \"K\": " << r.K << ", \"N\": " << r.N + << ", \"kv_heads\": " << r.kv + << ", \"op_mean_us\": " << json_num(r.mean_us) + << ", \"op_stdev_us\": " << json_num(r.stdev_us) + << ", \"kernel_mean_us\": " << json_num(r.kernel_us) + << ", \"kernel_median_us\": " << json_num(r.kernel_median_us) + << ", \"kernel_cov\": " << json_num(r.kernel_cov) + << ", \"gflops\": " << json_num(r.gflops) << ", \"dispatch\": \"" + << json_escape(r.dispatch) << "\"" << ", \"correctness\": \"" + << json_escape(r.correctness) << "\"" << ", \"detail\": \"" + << json_escape(r.detail_note) << "\"" + << ", \"ok\": " << (r.ok ? "true" : "false") << "}" + << (i + 1 < g_records.size() ? "," : "") << "\n"; + } + out << " ]\n}\n"; +} + void print_report(bool baseline_ran) { print_separator(); std::cout << "==================== RAW RESULTS ====================\n"; @@ -1723,6 +2025,11 @@ int main(int argc, char** argv) { bool linear = false, baseline = false, sdpa = false; bool correctness_only = false, sdpa_correctness_only = false; bool skip_correctness = false, list_only = false; + bool production_diff = false; + // Additive machine-readable output. Absent, every existing line is byte + // for byte what it was. + bool json_out = false; + std::string json_path; CaseFilter filter; for (int i = 1; i < argc; ++i) { const std::string arg = argv[i]; @@ -1736,10 +2043,17 @@ int main(int argc, char** argv) { correctness_only = true; } else if (arg == "--sdpa-correctness-only") { sdpa_correctness_only = true; + } else if (arg == "--production-diff") { + production_diff = true; } else if (arg == "--skip-correctness") { skip_correctness = true; } else if (arg == "--list") { list_only = true; + } else if (arg == "--json") { + json_out = true; + } else if (arg.rfind("--json-out=", 0) == 0) { + json_out = true; + json_path = arg.substr(11); } else if (arg.rfind("--model=", 0) == 0) { filter.model = arg.substr(8); } else if (arg.rfind("--scheme=", 0) == 0) { @@ -1802,7 +2116,8 @@ int main(int argc, char** argv) { << ",timestamp_period_ns=" << adapter->timestamp_period() << ",subgroup_size=" << adapter->subgroup_size() << ",coopmat=" << (adapter->supports_cooperative_matrix() ? "yes" : "no") - << "\n"; + << ",max_shared_mem_bytes=" + << adapter->max_compute_shared_memory_size() << "\n"; } print_separator(); @@ -1811,15 +2126,46 @@ int main(int argc, char** argv) { // Correctness gate: validates the tiled and coopmat linear kernels // (including the rank-3 dispatch check) before any perf time is spent. + auto finish_correctness = [&](bool ok) { + if (json_out) { + if (json_path.empty()) { + print_json_report(std::cout); + } else { + std::ofstream jf(json_path); + if (jf) { + print_json_report(jf); + } + } + } + return ok ? 0 : 1; + }; if (correctness_only) { - return run_linear_correctness(filter) ? 0 : 1; + return finish_correctness(run_linear_correctness(filter)); } if (sdpa_correctness_only) { - return run_sdpa_correctness() ? 0 : 1; + return finish_correctness(run_sdpa_correctness()); + } + if (production_diff) { + return finish_correctness(run_production_diff()); } if ((linear || baseline) && !skip_correctness) { if (!run_linear_correctness(filter)) { std::cout << "correctness gate FAILED -- not running the perf sweep\n"; + // Emit the JSON before returning. Without this the consumer sees no + // document at all and cannot distinguish a correctness failure from a + // crash -- which are different verdicts with different handling: one + // rejects the candidate, the other quarantines it and attempts device + // recovery. + if (json_out) { + if (json_path.empty()) { + print_json_report(std::cout); + } else { + std::ofstream jf(json_path); + if (jf) { + print_json_report(jf); + } + } + } return 1; } } @@ -1837,6 +2183,19 @@ int main(int argc, char** argv) { print_report(baseline); + if (json_out) { + if (json_path.empty()) { + print_json_report(std::cout); + } else { + std::ofstream jf(json_path); + if (!jf) { + std::cerr << "could not open " << json_path << " for --json-out\n"; + return 2; + } + print_json_report(jf); + } + } + // Exit code reflects dispatch sanity, not speed: every linear-suite // prefill buffer row must have dispatched coopmat, no coopmat may appear // where it can't (decode/forced-tiled/texture), nothing crashed, and From 009aa29dc7b4079d09a9c21b8ca42305cf116552 Mon Sep 17 00:00:00 2001 From: Yanwen Xu Date: Tue, 1 Sep 2026 22:27:58 -0700 Subject: [PATCH 15/28] [ET-VK] Promote coopMat-staged-A ("-tr" on zpg) as 8da4w coopmat default Replaces tsweep_dbuf4zpg_t128x64k32g42s32 with tsweep_dbuf4zpgtr_t128x64k32g42s32 as the shipped dq8ca coopmat linear default. Same tile, B-staging, and zp-hoist as dbuf4zpg; only the A-operand LDS staging changes, from a per-thread scalar scatter to a coopMatLoad(global)->coopMatStore(LDS) sequence. Validated on the sibling dq8ca-tr-staged-a-on-zpg branch (cut from this branch @ 1f3322ca22): -6.90%/-6.82%/-6.76% kern_us on 8B/3B/1B (46.50% -> 49.94% efficiency of int8 peak on 8B), 349.7 -> 364.5 tok/s real e2e prefill on 8B (+4.2%), vgpr_count 133 -> 128. Correctness: 10/10 consecutive clean (buffer) + 6/6 consecutive clean (texture3d) across all three model sizes, backed by an exhaustive host-side address-equivalence proof that the old and new A-staging schemes write the identical LDS address set. A follow-up tile re-sweep against this shader's own lower register pressure found no better tile. See openspec/changes/archive/2026-08-31-dq8ca-tr-staged-a-on-zpg and openspec/changes/coopmat-tr-tilesweep-4w-port for the full record. --- ...dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr.glsl | 752 ++++++++++++++++++ ...dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr.yaml | 60 ++ .../graph/ops/impl/QuantizedLinear.cpp | 53 +- 3 files changed, 859 insertions(+), 6 deletions(-) create mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr.glsl create mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr.glsl b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr.glsl new file mode 100644 index 00000000000..e9f49f99ef8 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr.glsl @@ -0,0 +1,752 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +/* + * "zpg" + "-tr" combination: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpg.glsl + * with its per-thread scalar A-staging replaced by + * linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr.glsl's coopMat-mediated A + * staging (coopMatLoad(global) -> coopmat<> -> coopMatStore(LDS)). This is + * an ADDITIVE combination, not a redesign: every non-A-staging block below + * (B staging: coalesced write, no skew; zp-hoist: izp/ifs applied once after + * the group loop via wcorr_sh; byte-parallel nibble widening; static + * A_ALWAYS_ACTIVE branch elision -- N/A here, see below; group epilog; + * bias/store epilogue) is byte-identical to dbuf4zpg's. Only the A-staging + * block (prologue load+store, main-loop prefetch+store) is dbuf4tr's, + * verbatim. + * + * dbuf4zpg's per-thread A staging used an `a_active` guard (statically always + * true when A_ACTIVE_THREADS == WG_SIZE, via the A_MAP_FULL-gated + * A_ALWAYS_ACTIVE macro). dbuf4tr's per-SUBGROUP tile map has no equivalent + * concept -- every subgroup participates via a `t < NUM_A_TILES` guard that + * depends only on gl_SubgroupID, not gl_LocalInvocationID.x -- so + * A_MAP_FULL/A_ALWAYS_ACTIVE is dropped entirely in this file; it would be + * dead code for the new A-staging block. + * + * Rationale for combining this way (not the reverse) and why this is worth + * building at all: see this change's design.md D0-D3. In short -- the only + * existing measurement of dbuf4tr's A-staging technique (28.72-30.51%, + * dq8ca-arch-redesign) was taken against dbuf4tr's own pre-zpg baseline + * (old B skew, no byte-parallel widening, no branch elision) -- a materially + * weaker shader than the 46.49-46.50% dbuf4zpg this file now combines it + * with. This file exists to answer whether that combination performs + * differently now that register pressure is already reduced. + * + * A staging (the actual delta from dbuf4zpg): + * dbuf4zpg: per-thread (m4, k4) ivec4 fetch, hoisted a_lds_off0/a_glb_row; + * only A_ACTIVE_THREADS invocations participate, each scattering + * 4 rows into Ash_int8 with 4 scalar stores. + * this file: per-SUBGROUP MMA_M x MMA_K tile fetch via coopMatLoad straight + * from a ROW-MAJOR (kPackedInt8_4W) int8 activation buffer, then + * coopMatStore into the same Ash_int8 slot -- dbuf4tr's mapping, + * unmodified (not re-derived; see design.md D3). + * + * t_packed_int8_input is therefore bound the same way dbuf4tr binds it: a + * SCALAR int8_t array in the kPackedInt8_4W layout (plain row-major int8, + * row stride K), produced by quantize_and_pack_4w_with_group_sums.glsl. + * QuantizedLinear.cpp's dq8ca_variant_wants_rowmajor_a() must recognize this + * file's variant token (tsweep_dbuf4zpgtr_t...) the same way it already + * recognizes tsweep_dbuf4tr_t/trm_t/trd_t, so graph-build time (packer + * selection) and dispatch time (kernel selection) cannot disagree. + * + * B CANNOT be coopmat-staged (int4 nibble unpack; a coopmat's per-lane layout + * is opaque to hand-assembly from unpacked registers) -- unchanged from both + * parent files. B staging below is dbuf4zpg's byte-parallel, coalesced, + * no-skew version, untouched. + * + * The loop structure is dbuf4's (both parents share it), unchanged: + * prologue: prefetch chunk 0 -> temp, store to slice 0 (no barrier) + * per iter: barrier -> prefetch(next) -> MMA(cur) -> store(next) + * kept nested (groups x chunks) with an unconditional group epilog -- + * flattening it crashes the Xclipse PAL compiler at large spec-resolved trip + * counts (see dbuf2's own header). + * + * Selected via + * ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4zpgtr_txkgs<32|64> + * (QuantizedLinear.cpp), additive to the tsweep_dbuf4zpg_t..., tsweep_dbuf4tr_t... + * and tsweep_t... namespaces. NOT the default -- unvalidated until it passes + * repeated test_llama_microbench --correctness-only runs (see + * dq8ca_coopmat_variant()'s comment on why a single pass is not proof). + * + * Performs: out[M,N] = dequant(int8_act) * dequant(int_w) (+ bias) + * via coopmat x coopmat -> coopmat on the matrix unit. + * + * Hard preconditions (dbuf4zpg's, plus dbuf4tr's row-major/alignment ones): + * M % WG_TILE_M == 0, N % WG_TILE_N == 0, K % WG_TILE_K == 0, + * group_size % WG_TILE_K == 0, K % 4 == 0, + * WG_TILE_M % MMA_M == 0, WG_TILE_K % MMA_K == 0, + * t_packed_int8_input in kPackedInt8_4W (row-major) layout, + * device exposes coopmatx-> at 16x16x16. + */ + +#version 450 core + +#extension GL_KHR_cooperative_matrix : require +#extension GL_KHR_memory_scope_semantics : require +#extension GL_KHR_shader_subgroup_basic : enable +#extension GL_EXT_shader_explicit_arithmetic_types : require +#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require +// 8-bit SSBO access: A is bound as a scalar int8_t array so that the +// coopMatLoad below has a MATCHING component type (see dbuf4tr's header for +// why the type must match on this driver). +#extension GL_EXT_shader_8bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_EXT_control_flow_attributes : enable + +#define PRECISION ${PRECISION} + +$if WEIGHT_NBITS == 4: + #define WEIGHT_INT4 + +$if HAS_BIAS: + #define HAS_BIAS + +$if WEIGHT_STORAGE == "buffer": + #define WEIGHT_BUFFER + +$if IO_STORAGE == "texture3d": + #define IO_TEXTURE + +layout(std430) buffer; + +#include "common.glslh" + +// Bindings — match add_linear_dqa_qw_node arg order: +// output(0), fp_input(1), packed_int8_input(2), int_input_sums(3 - unused), +// input_scales(4), input_zps(5), packed_weight(6), weight_sums(7), +// weight_scales(8), bias(9). +${layout_declare_tensor(B, "w", "t_output", "half", IO_STORAGE, is_scalar_array=True)} +// t_input is unread here -- the activations arrive already quantized in +// t_packed_int8_input -- but stays declared so the binding layout matches the +// dispatch site. It tracks IO_STORAGE so the two IO tensors stay consistent. +${layout_declare_tensor(B, "r", "t_input", "half", IO_STORAGE, is_scalar_array=False)} +// ROW-MAJOR (kPackedInt8_4W) packed activations, bound as a scalar int8_t +// array (row stride = K int8) -- dbuf4tr's binding, unchanged. The stock +// 4h4w layout dbuf4zpg uses is NOT row-major (component index selects a row, +// non-affine), so it cannot be addressed by any coopMatLoad. +${layout_declare_tensor(B, "r", "t_packed_int8_input", "int8", "buffer", is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_int8_input_sums", "int", "buffer", is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_int8_input_scales", "half", "texture3d")} +${layout_declare_tensor(B, "r", "t_int8_input_zps", "int8", "texture3d")} +${layout_declare_tensor(B, "r", "t_packed_weight", "int", WEIGHT_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_weight_sums", "int", "buffer", is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_weight_scales", "half", "buffer", is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_bias", "half", "buffer", is_scalar_array=True)} + +${layout_declare_ubo(B, "ivec4", "output_sizes")} +${layout_declare_ubo(B, "ivec4", "input_sizes")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +${layout_declare_spec_const(C, "int", "apply_bias", "0")} +// INT4 only; inert (0) for INT8 so the dispatcher's spec list lines up. +${layout_declare_spec_const(C, "int", "K4_per_group", "0")} +${layout_declare_spec_const(C, "int", "num_groups_arg", "0")} +${layout_declare_spec_const(C, "int", "out_N_arg", "0")} + +// Tile geometry +const uint MMA_M = ${MMA_M}; +const uint MMA_N = ${MMA_N}; +const uint MMA_K = ${MMA_K}; + +const uint WG_TILE_M = ${WG_TILE_M}; +const uint WG_TILE_N = ${WG_TILE_N}; +const uint WG_TILE_K = ${WG_TILE_K}; + +const uint SG_GRID_X = ${SG_GRID_X}; +const uint SG_GRID_Y = ${SG_GRID_Y}; +const uint SUBGROUP_SIZE = ${SUBGROUP_SIZE}; +const uint NUM_SUBGROUPS = SG_GRID_X * SG_GRID_Y; +const uint WG_SIZE = NUM_SUBGROUPS * SUBGROUP_SIZE; + +const uint SG_TILE_M = WG_TILE_M / SG_GRID_Y; +const uint SG_TILE_N = WG_TILE_N / SG_GRID_X; +const uint MMAS_PER_SG_M = SG_TILE_M / MMA_M; +const uint MMAS_PER_SG_N = SG_TILE_N / MMA_N; + +const uint A_SLAB_INT8 = WG_TILE_M * MMA_K; +const uint B_USEFUL_U32 = MMA_K / 4u; +// No skew + coalesced write -- dbuf4zpg's B fix, unchanged (this file does +// not touch B staging at all). +const uint B_STRIDE_U32 = B_USEFUL_U32; +const uint B_SLAB_U32 = WG_TILE_N * B_STRIDE_U32; +const uint NUM_K_SLABS = WG_TILE_K / MMA_K; + +const uint A_SLAB_U32 = A_SLAB_INT8 / 4u; +const uint A_STRIDE_U32 = MMA_K / 4u; + +// One ping-pong slice covers all K-slabs of one chunk. +const uint ASH_SLICE_U32 = NUM_K_SLABS * A_SLAB_U32; +const uint BSH_SLICE_U32 = NUM_K_SLABS * B_SLAB_U32; + +// Double-buffered MMA operand staging. +shared uint Ash_int8[2u * ASH_SLICE_U32]; +shared uint Bsh_int8[2u * BSH_SLICE_U32]; + +// Per-WG-tile-row activation params (loaded ONCE at WG start; constant +// across groups). +shared int izp_sh[WG_TILE_M]; // int32 (cast from int8 source) for broadcast +shared float ifs_sh[WG_TILE_M]; // float32 (cast from fp16 source) for broadcast + +// Per-(group, output-channel) weight params, ping-ponged by group parity. +// (For per-channel INT8 only slice 0 is ever used.) +shared float wsc_sh[2u * WG_TILE_N]; +// SUM_g wsc[g][n]*wsum[g][n] per output channel -- weight-side only, so it is +// accumulated once in the prologue. dbuf4zpg's zp-hoist, unchanged. +shared float wcorr_sh[WG_TILE_N]; + +#ifdef HAS_BIAS +shared float bias_sh[WG_TILE_N]; +#endif + +#ifdef IO_TEXTURE +// Result staging for the imageStore epilogue, mirroring the fp16 kernel: +// SG_GRID_Y bands of MMA_M rows, each WG_TILE_N wide, row-major. A full +// WG_TILE_M x WG_TILE_N buffer would cost SG_GRID_Y/MMAS_PER_SG_M x more LDS +// and wreck occupancy. float16_t-typed because coopMatStore needs it. +const uint CSH_ROWS = SG_GRID_Y * MMA_M; +shared float16_t Csh_out[CSH_ROWS * WG_TILE_N]; +#endif + +// Running fp32 accumulator (across all groups). +coopmat + result[MMAS_PER_SG_M][MMAS_PER_SG_N]; + +// Per-group int32 MMA accumulator. +coopmat + accum_int32[MMAS_PER_SG_M][MMAS_PER_SG_N]; + + +// Byte-parallel int4 -> int8 widening. dbuf4zpg's, unchanged (B-side only). +// +// The four nibbles this shader needs from one packed uint are ALREADY one per +// byte (bits 3:0 of each byte for parity 0, bits 7:4 for parity 1), so all four +// can be widened at once instead of with a per-nibble +// shift/mask/bias-subtract/mask chain. +// +// For v in [0,15] the biased value is v-8. `v ^ 8` is exactly the 4-bit +// two's-complement pattern of v-8, because -8 == +8 (mod 16): +// v=0 -> 0x8 -> -8 v=7 -> 0xF -> -1 +// v=8 -> 0x0 -> 0 v=15 -> 0x7 -> +7 +// so the only remaining work is sign-extending bit 3 into bits 7:4 per byte. +// `sgn * 0x1E` does that with no cross-byte carry: 0x08 * 0x1E == 0xF0 exactly, +// and sgn is at most 0x08080808 so the product is at most 0xF0F0F0F0. +// +// A naive `nib - 0x08080808` would NOT work -- it borrows across byte lanes +// whenever a nibble is < 8. Shifts must be on uint, not int, so the >> is +// logical rather than arithmetic. +// +// ~5 ops per 4 weights vs ~22 for the per-nibble chain; bit-identical output. +uint widen_nibbles(const uint w, const uint parity) { + const uint nib = (parity == 0u) ? (w & 0x0F0F0F0Fu) : ((w >> 4u) & 0x0F0F0F0Fu); + const uint p = nib ^ 0x08080808u; + const uint sgn = p & 0x08080808u; + return p | (sgn * 0x1Eu); +} + +void main() { + const uvec2 tileID = uvec2(gl_WorkGroupID.xy); + const uvec2 warpInTile = uvec2( + gl_SubgroupID % SG_GRID_X, + gl_SubgroupID / SG_GRID_X); + + const uint K = uint(input_sizes.x); + const uint N = uint(output_sizes.x); + const uint N4 = (N + 3u) / 4u; + const uint nblocks_x_A = (K + 3u) >> 2u; + // A row stride in INT8 elements (dbuf4tr's binding is row-major int8, not + // the 4h4w ivec4 block layout dbuf4zpg used -- so A addressing below is in + // int8 elements, not int, and derived from nblocks_x_A so it matches the + // packer's `m_row * K4 + k4` addressing exactly (K % 4 == 0 makes them + // equal to K directly). + const uint a_row_stride_i8 = nblocks_x_A * 4u; + +#ifdef WEIGHT_INT4 + const uint num_groups = uint(num_groups_arg); + const uint CHUNKS_PER_GROUP = uint(K4_per_group) * 4u / WG_TILE_K; +#else + // Per-channel: a single quant "group" spanning all of K. The nested + // groups x chunks loop below collapses to a flat chunk loop, the wsum/wsc + // ping-pong never crosses a boundary, and the epilog runs exactly once. + const uint num_groups = 1u; + const uint CHUNKS_PER_GROUP = uint(num_groups_arg); +#endif + const uint num_chunks = num_groups * CHUNKS_PER_GROUP; + + const uint tile_m_start = WG_TILE_M * tileID.y; + const uint tile_n_start = WG_TILE_N * tileID.x; + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + result[i][j] = coopmat(0.0); + accum_int32[i][j] = coopmat(0); + } + } + + const uint K_BLOCKS_PER_CHUNK = WG_TILE_K >> 2u; + + // --- A staging tile map (dbuf4tr's, unmodified): one MMA_M x MMA_K coopmat + // tile per subgroup per slot, dealt round-robin across the + // NUM_SUBGROUPS subgroups so every subgroup participates. Replaces + // dbuf4zpg's per-thread (m4, k4) map / a_active guard entirely -- see + // design.md D3 for why this is reused as-is, not re-derived. + const uint A_TILES_M = WG_TILE_M / MMA_M; + const uint A_TILES_K = WG_TILE_K / MMA_K; // == NUM_K_SLABS + const uint NUM_A_TILES = A_TILES_M * A_TILES_K; + const uint A_TILES_PER_SG = (NUM_A_TILES + NUM_SUBGROUPS - 1u) / NUM_SUBGROUPS; + +#ifdef WEIGHT_INT4 + // --- B staging thread map: (block, col) slots; each slot extracts one + // ColumnMajor LDS uint (4 K-contiguous sign-extended int8) --- + const uint B_TOTAL_SLOTS = K_BLOCKS_PER_CHUNK * WG_TILE_N; + const uint B_SLOTS_PER_THREAD = B_TOTAL_SLOTS / WG_SIZE; + const uint N8_PER_TILE = WG_TILE_N >> 3u; +#else + // --- B staging thread map: one (k4, n4) ivec4 block per active thread --- + const uint B_FETCH_SLOTS = K_BLOCKS_PER_CHUNK * (WG_TILE_N >> 2u); + const uint N4_PER_TILE = WG_TILE_N >> 2u; + const uint b_k4_in_chunk = gl_LocalInvocationID.x / N4_PER_TILE; + const uint b_n_uint_col = gl_LocalInvocationID.x % N4_PER_TILE; + const bool b_active = gl_LocalInvocationID.x < B_FETCH_SLOTS; +#endif + + // ===== INTERVENTION F: hoist loop-invariant B staging index math ===== + // dbuf4zpg's, unchanged -- B staging is untouched by this file's A-staging + // swap. See dbuf4zpg's header for the full rationale (ablation-attributed + // -16.8% block, mostly index arithmetic not memory traffic). +#ifdef WEIGHT_INT4 + uint b_lds_off[B_SLOTS_PER_THREAD]; // LDS store offset within a slice + uint b_comp[B_SLOTS_PER_THREAD]; // which ivec4 component feeds this slot + uint b_par[B_SLOTS_PER_THREAD]; // nibble parity for this slot + uint b_n8blk[B_SLOTS_PER_THREAD]; // global texel column (N/8 blocks) + uint b_k4off[B_SLOTS_PER_THREAD]; // k4 offset of this slot within a chunk + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { + const uint a = gl_LocalInvocationID.x + si * WG_SIZE; + const uint slab_idx = a / B_SLAB_U32; + const uint local_a = a % B_SLAB_U32; + const uint n_col = local_a / B_STRIDE_U32; + const uint k4_in_slab = local_a % B_STRIDE_U32; + const uint k4_in_chunk = slab_idx * (MMA_K >> 2u) + k4_in_slab; + const uint n8_in_tile = n_col >> 3u; + const uint rem = n_col & 7u; + b_lds_off[si] = a; + b_comp[si] = rem & 3u; + b_par[si] = rem >> 2u; + b_n8blk[si] = (tile_n_start >> 3u) + n8_in_tile; + b_k4off[si] = k4_in_chunk; + } +#endif + + // Prefetch temp registers. temp_A is a coopmat array (dbuf4tr's A-staging + // technique); indices into it are [[unroll]]-resolved compile-time + // constants, never dynamic -- dynamic indexing of a coopmat array is + // exactly the construct the Xclipse/AMD-PAL compiler has miscompiled + // before. + coopmat + temp_A[A_TILES_PER_SG]; +#ifdef WEIGHT_INT4 + ivec4 temp_B[B_SLOTS_PER_THREAD]; + float temp_wsc; +#else + ivec4 temp_B; +#endif + + // ========================================================= + // PROLOGUE + // ========================================================= + if (gl_LocalInvocationID.x < (WG_TILE_M >> 2u)) { + const uint m4 = (tile_m_start >> 2u) + gl_LocalInvocationID.x; + const vec4 sc = vec4(texelFetch(t_int8_input_scales, ivec3(m4, 0, 0), 0)); + const ivec4 zp = texelFetch(t_int8_input_zps, ivec3(m4, 0, 0), 0); + const uint base = gl_LocalInvocationID.x * 4u; + ifs_sh[base + 0u] = sc.x; ifs_sh[base + 1u] = sc.y; + ifs_sh[base + 2u] = sc.z; ifs_sh[base + 3u] = sc.w; + izp_sh[base + 0u] = zp.x; izp_sh[base + 1u] = zp.y; + izp_sh[base + 2u] = zp.z; izp_sh[base + 3u] = zp.w; + } + // Group 0 weight scales -> slice 0, and the hoisted weight-side correction + // SUM_g wsc[g][n]*wsum[g][n] accumulated across ALL groups. dbuf4zpg's + // zp-hoist, unchanged. + if (gl_LocalInvocationID.x < WG_TILE_N) { + const uint n_idx = tile_n_start + gl_LocalInvocationID.x; + f16vec4 sv0 = t_weight_scales[n_idx >> 2u]; + wsc_sh[gl_LocalInvocationID.x] = float(sv0[n_idx & 3u]); + + float corr = 0.0; + for (uint g = 0; g < num_groups; ++g) { + f16vec4 sv = t_weight_scales[g * N4 + (n_idx >> 2u)]; + corr += float(sv[n_idx & 3u]) * float(t_weight_sums[g * N + n_idx]); + } + wcorr_sh[gl_LocalInvocationID.x] = corr; + } + memoryBarrierShared(); + barrier(); + + // NOTE: dbuf4zpg builds izp_bcast/ifs_bcast AFTER the group loop, not here + // -- that is the register-pressure saving zp-hoist buys. Unchanged. + + // dbuf4: prefetch chunk 0 into temp registers, THEN store to slice 0 (no + // barrier here -- the main loop's first iteration barriers before + // reading slice 0). + // + // A staging (dbuf4tr's technique): per-subgroup coopMatLoad straight from + // the row-major global buffer. + [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { + const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; + if (t < NUM_A_TILES) { + const uint tm = t / A_TILES_K; + const uint tk = t % A_TILES_K; + coopMatLoad( + temp_A[s], t_packed_int8_input, + (tile_m_start + tm * MMA_M) * a_row_stride_i8 + tk * MMA_K, + a_row_stride_i8, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#ifdef WEIGHT_INT4 + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { +#ifdef WEIGHT_BUFFER + temp_B[si] = t_packed_weight[(b_n8blk[si] * nblocks_x_A) + b_k4off[si]]; +#else + temp_B[si] = texelFetch(t_packed_weight, ivec2(b_k4off[si], b_n8blk[si]), 0); +#endif + } +#else + if (b_active) { + const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; +#ifdef WEIGHT_BUFFER + temp_B = t_packed_weight[(b_k4_in_chunk * N4) + block_x_w]; +#else + temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, b_k4_in_chunk), 0); +#endif + } +#endif + { + // store chunk 0 -> slice 0 + // A staging (dbuf4tr's technique): coopMatStore into the same Ash_int8 + // slot layout dbuf4zpg's scalar scatter used to write. + [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { + const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; + if (t < NUM_A_TILES) { + const uint tm = t / A_TILES_K; + const uint tk = t % A_TILES_K; + coopMatStore( + temp_A[s], Ash_int8, + tk * A_SLAB_U32 + (tm * MMA_M) * A_STRIDE_U32, + A_STRIDE_U32, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#ifdef WEIGHT_INT4 + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { + Bsh_int8[b_lds_off[si]] = + widen_nibbles(uint(temp_B[si][b_comp[si]]), b_par[si]); + } +#else + if (b_active) { + const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); + const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); + const uint n_col_base = b_n_uint_col * 4u; + [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { + Bsh_int8[slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab] = + uint(temp_B[n_in_blk]); + } + } +#endif + } + + // ========================================================= + // MAIN LOOP (dbuf4) — nested groups x chunks (kept nested; flattening it + // with a conditional coopmat epilog crashes the Xclipse PAL compiler at + // large spec-resolved trip counts). One barrier per chunk. Chunk + // iteration (global index `chunk`): + // 1. barrier — A/B slice (chunk%2) fully written; on the first chunk + // of group g, wsc slice (g%2) is too. + // 2. prefetch — chunk+1 (A tiles, B blocks) into temp; when chunk+1 + // starts a new group, also its wsc element. Skipped + // entirely on the final chunk. + // 3. int8 MMA — on slice (chunk%2) into accum_int32. + // 4. store — temp -> A/B slice ((chunk+1)%2), unpacking the weight; + // on a group boundary, wsc -> slice ((g+1)%2). + // The group epilog runs unconditionally at the tail of each group. + // ========================================================= + uint chunk = 0; + for (uint group_i = 0; group_i < num_groups; ++group_i) { + for (uint inner = 0; inner < CHUNKS_PER_GROUP; ++inner, ++chunk) { + const bool has_next = chunk + 1u < num_chunks; + const bool group_crossing = has_next && (inner + 1u == CHUNKS_PER_GROUP); + const uint cur_a = (chunk % 2u) * ASH_SLICE_U32; + const uint cur_b = (chunk % 2u) * BSH_SLICE_U32; + const uint nxt_a = ((chunk + 1u) % 2u) * ASH_SLICE_U32; + const uint nxt_b = ((chunk + 1u) % 2u) * BSH_SLICE_U32; + + barrier(); + + // --- 2. prefetch chunk+1 -> temp --- + if (has_next) { + const uint chunkK_nxt = (chunk + 1u) * WG_TILE_K; + // A staging (dbuf4tr's technique): coopMatLoad straight from global. + [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { + const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; + if (t < NUM_A_TILES) { + const uint tm = t / A_TILES_K; + const uint tk = t % A_TILES_K; + coopMatLoad( + temp_A[s], t_packed_int8_input, + (tile_m_start + tm * MMA_M) * a_row_stride_i8 + chunkK_nxt + + tk * MMA_K, + a_row_stride_i8, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#ifdef WEIGHT_INT4 + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { + const uint k4_blk = (chunkK_nxt >> 2u) + b_k4off[si]; +#ifdef WEIGHT_BUFFER + temp_B[si] = t_packed_weight[(b_n8blk[si] * nblocks_x_A) + k4_blk]; +#else + temp_B[si] = texelFetch(t_packed_weight, ivec2(k4_blk, b_n8blk[si]), 0); +#endif + } + if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { + const uint n_idx = tile_n_start + gl_LocalInvocationID.x; + f16vec4 sv = t_weight_scales[(group_i + 1u) * N4 + (n_idx >> 2u)]; + temp_wsc = float(sv[n_idx & 3u]); + } +#else + if (b_active) { + const uint block_y_w = (chunkK_nxt >> 2u) + b_k4_in_chunk; + const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; +#ifdef WEIGHT_BUFFER + temp_B = t_packed_weight[(block_y_w * N4) + block_x_w]; +#else + temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, block_y_w), 0); +#endif + } +#endif + } + + // --- 3. int8 MMA on the cur slice --- + [[unroll]] for (uint k = 0; k < NUM_K_SLABS; ++k) { + const uint slab_a_base_u32 = cur_a + k * A_SLAB_U32; + const uint slab_b_base_u32 = cur_b + k * B_SLAB_U32; + + coopmat matA[MMAS_PER_SG_M]; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + const uint row_a = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + coopMatLoad( + matA[i], Ash_int8, + slab_a_base_u32 + row_a * A_STRIDE_U32, + A_STRIDE_U32, + gl_CooperativeMatrixLayoutRowMajor); + } + + coopmat matB; + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint col_b = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopMatLoad( + matB, Bsh_int8, + slab_b_base_u32 + col_b * B_STRIDE_U32, + B_STRIDE_U32, + gl_CooperativeMatrixLayoutColumnMajor); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + accum_int32[i][j] = coopMatMulAdd(matA[i], matB, accum_int32[i][j]); + } + } + } + + // --- 4. store temp (chunk+1) -> nxt slice --- + if (has_next) { + // A staging (dbuf4tr's technique): coopMatStore into the nxt slice. + [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { + const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; + if (t < NUM_A_TILES) { + const uint tm = t / A_TILES_K; + const uint tk = t % A_TILES_K; + coopMatStore( + temp_A[s], Ash_int8, + nxt_a + tk * A_SLAB_U32 + (tm * MMA_M) * A_STRIDE_U32, + A_STRIDE_U32, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#ifdef WEIGHT_INT4 + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { + Bsh_int8[nxt_b + b_lds_off[si]] = + widen_nibbles(uint(temp_B[si][b_comp[si]]), b_par[si]); + } + if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { + const uint wbase_nxt = ((group_i + 1u) % 2u) * WG_TILE_N; + wsc_sh[wbase_nxt + gl_LocalInvocationID.x] = temp_wsc; + } +#else + if (b_active) { + const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); + const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); + const uint n_col_base = b_n_uint_col * 4u; + [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { + Bsh_int8[nxt_b + slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab] = + uint(temp_B[n_in_blk]); + } + } +#endif + } + } // chunks + + // --- Group epilog: scale-only accumulate, reset accum --- + // dbuf4zpg's, unchanged. Just result += float(acc) * wsc. The + // zero-point subtract and the ifs multiply are hoisted out of the group + // loop (applied once below). + { + const uint wbase = (group_i % 2u) * WG_TILE_N; + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + + coopmat wsc_bcast; + coopMatLoad( + wsc_bcast, wsc_sh, + wbase + local_n_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutRowMajor); + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + result[i][j] += + coopmat( + accum_int32[i][j]) * wsc_bcast; + accum_int32[i][j] = coopmat(0); + } + } + } + } // groups + + // --- Hoisted correction, applied ONCE: -------------------------------- + // result = ifs * ( result - izp * SUM_g wsc_g*wsum_g ) + // dbuf4zpg's, unchanged. izp/ifs are loaded here rather than before the + // group loop so they are not live across it. + { + coopmat + izpf_bcast[MMAS_PER_SG_M]; + coopmat + ifs_bcast[MMAS_PER_SG_M]; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + const uint local_m_base = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + coopmat izp_i; + coopMatLoad( + izp_i, izp_sh, local_m_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutColumnMajor); + izpf_bcast[i] = + coopmat(izp_i); + coopMatLoad( + ifs_bcast[i], ifs_sh, local_m_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutColumnMajor); + } + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopmat wcorr_bcast; + coopMatLoad( + wcorr_bcast, wcorr_sh, local_n_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutRowMajor); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + result[i][j] = ifs_bcast[i] * (result[i][j] - izpf_bcast[i] * wcorr_bcast); + } + } + } + + // --- Bias (optional) --- +#ifdef HAS_BIAS + if (apply_bias > 0) { + for (uint t = gl_LocalInvocationID.x; t < WG_TILE_N; t += WG_SIZE) { + bias_sh[t] = float(t_bias[tile_n_start + t]); + } + memoryBarrierShared(); + barrier(); + } +#endif + + // --- Store result tile --- + // N for the store address math MUST come from the spec constant, not the + // sizes UBO (see out_N_arg above). +#ifdef IO_TEXTURE + // Epilogue iteration i drains accumulator row-block i from EVERY subgroup + // into Csh_out at once, so the SG_GRID_Y bands it holds are disjoint global + // row ranges; the whole workgroup then imageStores them. lr / MMA_M is the + // writing subgroup's warpInTile.y, so the global row reproduces the buffer + // path's gi exactly. + // + // PORTABILITY NOTE: the barrier() in the loop body keeps this loop rolled + // despite [[unroll]], so result[i][j] IS dynamically indexed. Coopmat arrays + // are opaque per-lane storage and dynamic indexing is exactly the construct + // the Xclipse/AMD-PAL compiler has broken before -- check this first if the + // texture variants miscompile on M51. + const uint CSH_TEXELS_PER_ROW = WG_TILE_N / 4u; + const uint CSH_TEXELS = CSH_ROWS * CSH_TEXELS_PER_ROW; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + // Guards Csh_out against the previous iteration's readers. Inert on i == 0 + // but must stay unconditional to remain workgroup-uniform. + barrier(); + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { +#ifdef HAS_BIAS + if (apply_bias > 0) { + const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopmat bias_tile; + coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); + result[i][j] += bias_tile; + } +#endif + coopmat out_tile = + coopmat(result[i][j]); + coopMatStore( + out_tile, Csh_out, + warpInTile.y * MMA_M * WG_TILE_N + + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j), + WG_TILE_N, + gl_CooperativeMatrixLayoutRowMajor); + } + memoryBarrierShared(); + barrier(); + + for (uint t = gl_LocalInvocationID.x; t < CSH_TEXELS; t += WG_SIZE) { + const uint lr = t / CSH_TEXELS_PER_ROW; + const uint lc4 = t % CSH_TEXELS_PER_ROW; + const uint m = + tile_m_start + (lr / MMA_M) * SG_TILE_M + i * MMA_M + (lr % MMA_M); + const uint base = lr * WG_TILE_N + lc4 * 4u; + imageStore( + t_output, + ivec3(tile_n_start / 4u + lc4, m, 0), + vec4( + float(Csh_out[base]), + float(Csh_out[base + 1u]), + float(Csh_out[base + 2u]), + float(Csh_out[base + 3u]))); + } + } +#else + const uint N_out = uint(out_N_arg); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint gi = tile_m_start + MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + const uint gj = tile_n_start + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + +#ifdef HAS_BIAS + if (apply_bias > 0) { + const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopmat bias_tile; + coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); + result[i][j] += bias_tile; + } +#endif + + coopmat out_tile = + coopmat(result[i][j]); + coopMatStore( + out_tile, t_output, + gi * N_out + gj, N_out, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#endif // IO_TEXTURE +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr.yaml b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr.yaml new file mode 100644 index 00000000000..9f5a6af571b --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr.yaml @@ -0,0 +1,60 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# "zpg" + "-tr" combination: tsweep_dbuf4zpg with its per-thread scalar +# A-staging replaced by tsweep_dbuf4tr's coopMat-mediated A staging. B +# staging, zp-hoist, byte-parallel nibble widening are all dbuf4zpg's, +# unchanged -- only A-staging differs. Requires t_packed_int8_input in the +# ROW-MAJOR kPackedInt8_4W layout (same requirement as tsweep_dbuf4tr). +# +# PROMOTED 2026-09-01 as the shipped default -- see QuantizedLinear.cpp's +# dq8ca_coopmat_variant() for the full validation record (10/10 buffer, 6/6 +# texture3d correctness across 1B/3B/8B; +4.2% real e2e prefill; 46.50% -> +# 49.94% efficiency on 8B). Also selectable explicitly via +# ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4zpgtr_txkgs. +# See openspec/changes/archive/2026-08-31-dq8ca-tr-staged-a-on-zpg. +# +# Single seed tile: the prior dbuf4zpg default's own tile (t128x64k32g42s32). +# A re-sweep against this shader's own (lower) register-pressure profile was +# run as a follow-up (coopmat-tr-tilesweep-4w-port) and found no better tile +# -- this remains the best known geometry. + +linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr: + parameter_names_with_default_values: + PRECISION: highp + HAS_BIAS: false + IO_STORAGE: buffer + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + MMA_M: 16 + MMA_N: 16 + MMA_K: 16 + WG_TILE_M: 64 + WG_TILE_N: 32 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + shader_variants: + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr_t128x64k32g42s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr_t128x64k32g42s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp index a1a773cac1b..3317b207fdf 100644 --- a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp @@ -115,9 +115,14 @@ static const char* const kDq8caTsweepPrefixes[] = { // zpi + compile-time elision of the statically-true a_active guard // (intervention G of dq8ca-prefill-stall-reduction), combined with the // dbuf4 default's own B_STRIDE_U32 skew removal + coalesced B-store - // rewrite. PROMOTED 2026-08-28 as the shipped default -- see - // dq8ca_coopmat_variant() below. + // rewrite. Superseded 2026-09-01 by tsweep_dbuf4zpgtr_t below -- kept + // listed (env-var-selectable) for comparison/rollback. "tsweep_dbuf4zpg_t", + // dbuf4zpg with its per-thread scalar A-staging replaced by a + // coopMat-mediated coopMatLoad(global)->coopMatStore(LDS) sequence (B + // staging/zp-hoist/nibble-widening unchanged). PROMOTED 2026-09-01 as + // the shipped default -- see dq8ca_coopmat_variant() below. + "tsweep_dbuf4zpgtr_t", // (dq8ca-dequant-unpack-ablation Addendum 11 -- abl_aconst/abl_areadc/ // abl_abconst -- were measurement-only variants deleted once each // attribution was recorded; see openspec/changes/dq8ca-dequant-unpack- @@ -210,8 +215,43 @@ static const std::string& q4gsw_coopmat_variant() { static const std::string& dq8ca_coopmat_variant() { // Default (no ET_VK_DQ8CA_COOPMAT_VARIANT set): - // tsweep_dbuf4zpg_t128x64k32g42s32 (WG_TILE 128x64x32, SG_GRID 4x2, wave32) - // -- PROMOTED 2026-08-28 from the prior default + // tsweep_dbuf4zpgtr_t128x64k32g42s32 (WG_TILE 128x64x32, SG_GRID 4x2, + // wave32) -- PROMOTED 2026-09-01 from tsweep_dbuf4zpg_t128x64k32g42s32 + // (below). Same tile/B-staging/zp-hoist as dbuf4zpg; the only change is + // A-staging: per-thread scalar scatter into Ash_int8 replaced by a + // coopMatLoad(global)->coopmat<>->coopMatStore(LDS) sequence (requires + // t_packed_int8_input in the ROW-MAJOR kPackedInt8_4W layout -- see + // dq8ca_variant_wants_rowmajor_a() above). + // + // Real, on-device, correctness-validated on the sibling + // `dq8ca-tr-staged-a-on-zpg` branch (cut from this branch @ 1f3322ca22): + // microbench kern_us -6.90%/-6.82%/-6.76% on 8B/3B/1B respectively (46.50% + // -> 49.94% efficiency of int8 peak on 8B), agreeing to within 0.14pp + // across all three model sizes with spreads 2-3 orders of magnitude + // smaller than the delta. Real e2e prefill (`ET_VK_TEXTURE_COOPMAT=1 + // ET_VK_EXECUTE_NODE_THRESHOLD=32`, xgpusw-debug08/00000b750f413c33, + // maxpin 980/5333/934): 349.7 -> 364.5 tok/s on 8B, +4.2%. vgpr_count fell + // 133 -> 128. Correctness: 10/10 consecutive clean (buffer) + 6/6 + // consecutive clean (texture3d) across all three model sizes, dispatch + // confirmed by kernel name each time; also backed by an exhaustive + // host-side address-equivalence proof that the old per-thread and new + // per-subgroup A-staging schemes write the identical Ash_int8 address set + // (test_llama_microbench's bench_reference can't cover M/N>256, so this + // substitutes for an M=2048 output differential -- same substitution + // dq8ca-prefill-stall-reduction used for the same harness limitation). + // A follow-up tile re-sweep against this shader's own (lower) + // register-pressure profile (coopmat-tr-tilesweep-4w-port, 2026-09-01) + // tried 7 further candidates (including the two 256-row shapes structurally + // undispatchable under this harness's fixed-M=128 rank-3 case) and found no + // better tile; deep PAL-counter/ISA/ablation follow-up work in that same + // change found no further lever either (both A's and B's remaining costs + // are small and near their practical floor) -- t128x64k32g42s32 is the + // best known geometry for this shader. + // Full record: openspec/changes/archive/2026-08-31-dq8ca-tr-staged-a-on-zpg + // and openspec/changes/coopmat-tr-tilesweep-4w-port. + // + // Prior default (dbuf4zpg, still selectable via + // tsweep_dbuf4zpg_t128x64k32g42s32), PROMOTED 2026-08-28 from // tsweep_dbuf4_t128x128k64g81s64 (below). Combines every validated 8da4w win // to date: zp-hoist (zero-point // + activation-scale correction moved out of the per-group loop), the @@ -288,7 +328,7 @@ static const std::string& dq8ca_coopmat_variant() { static const std::string variant = [] { const char* env = std::getenv("ET_VK_DQ8CA_COOPMAT_VARIANT"); if (!env) { - return std::string("tsweep_dbuf4zpg_t128x64k32g42s32"); + return std::string("tsweep_dbuf4zpgtr_t128x64k32g42s32"); } const std::string v(env); if (is_dq8ca_shippable_token(v)) { @@ -545,7 +585,8 @@ static bool dq8ca_variant_wants_rowmajor_a() { const std::string& v = dq8ca_coopmat_variant(); return v.rfind("tsweep_dbuf4tr_t", 0) == 0 || v.rfind("tsweep_dbuf4trm_t", 0) == 0 || - v.rfind("tsweep_dbuf4trd_t", 0) == 0; + v.rfind("tsweep_dbuf4trd_t", 0) == 0 || + v.rfind("tsweep_dbuf4zpgtr_t", 0) == 0; } // Mirrors the coopmat branch of pick_linear_dqa_qw_shader() so graph-build time From 7892151a81c0156f8ce48a96c6872522e62d30cf Mon Sep 17 00:00:00 2001 From: Yanwen Xu Date: Tue, 1 Sep 2026 22:27:58 -0700 Subject: [PATCH 16/28] [ET-VK] Fence the LDS staging barriers in both shipped coopmat linear defaults barrier() alone does NOT order shared-memory stores against a subsequent coopMatLoad on the M51 Xclipse/AMD-PAL driver; it needs an explicit memoryBarrierShared(). The symptom is recognisable and quiet: one stale MMA_M-row band of the A operand (16 consecutive rows, all WG_TILE_N columns), silently wrong output, no crash and no DEVICE_LOST, on roughly 2.5% of runs. Found 2026-09-02 in sdpa_compute_out_coopmat.glsl. Both shipped coopmat linear defaults had this defect on their double-buffer staging path, as did the 8da4w env-var fallback. The memoryBarrierShared() pairs already present in these files all guarded wcorr_sh / bias_sh / Csh_out -- none guarded Ash_int8/Bsh_int8, i.e. none guarded the store -> coopMatLoad ordering point that runs every chunk. linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr.glsl (8da4w default) +2 linear_q4gsw_coopmat_tsweep_dbuf4.glsl (4w default) +3 linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpg.glsl (8da4w fallback) +2 All three now have 0 bare barrier() calls. The diff is exactly 7 memoryBarrierShared() lines plus comments; QuantizedLinear.cpp is comment-only. COST: NONE, measured, not assumed. Interleaved same-binary A/B (unfenced clones of both defaults compiled alongside the fixed ones, alternating round-by-round so drift hits both arms equally), 8B prefill, buffer, 4 rounds each: 8da4w fenced 70625.1 us vs unfenced 70624.6 us -> +0.001% 4w fenced 95740.4 us vs unfenced 95741.9 us -> -0.002% Observed spread across all 8 runs was 0.037% (8da4w) and 0.017% (4w), and the paired per-round deltas straddle zero in both cases, so the fence cost is below the resolution of the measurement. The fence is NOT optimised away -- verified in the compiled SPIR-V rather than inferred from the timing: OpMemoryBarrier goes 1 -> 2 (8da4w) and 0 -> 2 (4w) between the unfenced and fenced builds, with OpControlBarrier unchanged at 2. Without that check a "no cost" result would be consistent with the compiler having stripped the fence, which would also mean it fixes nothing. Correctness re-gated after the driver reflash noted below: 14 PASSED / 0 FAILED for both defaults, correct shader confirmed dispatched. Measured on xgpusw-debug08 / 00000bf70c579c33, driver md5 1eea300aa3974ff8974d04808c9ff394 (SUMD main@fafb46ae9c0d, re-verified after a reflash -- a crash-to-bootloader recovery on this board silently reverted /vendor to the BSP factory driver), clocks pinned 980/5333/934. --- ...r_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpg.glsl | 14 +++++++++++++ ...dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr.glsl | 14 +++++++++++++ .../linear_q4gsw_coopmat_tsweep_dbuf4.glsl | 21 +++++++++++++++++++ .../graph/ops/impl/QuantizedLinear.cpp | 5 +++++ 4 files changed, 54 insertions(+) diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpg.glsl b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpg.glsl index 449d1710722..d5af04c92c7 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpg.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpg.glsl @@ -492,6 +492,13 @@ void main() { const uint nxt_a = ((chunk + 1u) % 2u) * ASH_SLICE_U32; const uint nxt_b = ((chunk + 1u) % 2u) * BSH_SLICE_U32; + // coopmat-lds-fence: barrier() alone does NOT order shared stores against a + // subsequent coopMatLoad on the M51 Xclipse/AMD-PAL driver -- symptom is one + // stale MMA_M-row band of A, all columns, ~2.5% of runs, no crash (observed + // 2026-09-02 in sdpa_compute_out_coopmat.glsl). Measured cost of the fence: + // none (see this change's results). See memory + // `coopmat-lds-needs-explicit-memorybarriershared`. + memoryBarrierShared(); barrier(); // --- 2. prefetch chunk+1 -> temp --- @@ -680,6 +687,13 @@ void main() { [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { // Guards Csh_out against the previous iteration's readers. Inert on i == 0 // but must stay unconditional to remain workgroup-uniform. + // coopmat-lds-fence: barrier() alone does NOT order shared stores against a + // subsequent coopMatLoad on the M51 Xclipse/AMD-PAL driver -- symptom is one + // stale MMA_M-row band of A, all columns, ~2.5% of runs, no crash (observed + // 2026-09-02 in sdpa_compute_out_coopmat.glsl). Measured cost of the fence: + // none (see this change's results). See memory + // `coopmat-lds-needs-explicit-memorybarriershared`. + memoryBarrierShared(); barrier(); [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { #ifdef HAS_BIAS diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr.glsl b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr.glsl index e9f49f99ef8..97eed7ddf67 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr.glsl @@ -484,6 +484,13 @@ void main() { const uint nxt_a = ((chunk + 1u) % 2u) * ASH_SLICE_U32; const uint nxt_b = ((chunk + 1u) % 2u) * BSH_SLICE_U32; + // coopmat-lds-fence: barrier() alone does NOT order shared stores against a + // subsequent coopMatLoad on the M51 Xclipse/AMD-PAL driver -- symptom is one + // stale MMA_M-row band of A, all columns, ~2.5% of runs, no crash (observed + // 2026-09-02 in sdpa_compute_out_coopmat.glsl). Measured cost of the fence: + // none (see this change's results). See memory + // `coopmat-lds-needs-explicit-memorybarriershared`. + memoryBarrierShared(); barrier(); // --- 2. prefetch chunk+1 -> temp --- @@ -686,6 +693,13 @@ void main() { [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { // Guards Csh_out against the previous iteration's readers. Inert on i == 0 // but must stay unconditional to remain workgroup-uniform. + // coopmat-lds-fence: barrier() alone does NOT order shared stores against a + // subsequent coopMatLoad on the M51 Xclipse/AMD-PAL driver -- symptom is one + // stale MMA_M-row band of A, all columns, ~2.5% of runs, no crash (observed + // 2026-09-02 in sdpa_compute_out_coopmat.glsl). Measured cost of the fence: + // none (see this change's results). See memory + // `coopmat-lds-needs-explicit-memorybarriershared`. + memoryBarrierShared(); barrier(); [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { #ifdef HAS_BIAS diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_coopmat_tsweep_dbuf4.glsl b/backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_coopmat_tsweep_dbuf4.glsl index 033afc969e8..08447b06059 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_coopmat_tsweep_dbuf4.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_coopmat_tsweep_dbuf4.glsl @@ -265,6 +265,13 @@ void main() { const uint nxt_base_A = ((chunk + 1u) % 2u) * ASH_SLICE; const uint nxt_base_B = ((chunk + 1u) % 2u) * BSH_SLICE; + // coopmat-lds-fence: barrier() alone does NOT order shared stores against a + // subsequent coopMatLoad on the M51 Xclipse/AMD-PAL driver -- symptom is one + // stale MMA_M-row band of A, all columns, ~2.5% of runs, no crash (observed + // 2026-09-02 in sdpa_compute_out_coopmat.glsl). Measured cost of the fence: + // none (see this change's results). See memory + // `coopmat-lds-needs-explicit-memorybarriershared`. + memoryBarrierShared(); barrier(); // --- prefetch chunk+1 -> temp --- @@ -339,6 +346,13 @@ void main() { const uint cur_base_A = (chunk % 2u) * ASH_SLICE; const uint cur_base_B = (chunk % 2u) * BSH_SLICE; + // coopmat-lds-fence: barrier() alone does NOT order shared stores against a + // subsequent coopMatLoad on the M51 Xclipse/AMD-PAL driver -- symptom is one + // stale MMA_M-row band of A, all columns, ~2.5% of runs, no crash (observed + // 2026-09-02 in sdpa_compute_out_coopmat.glsl). Measured cost of the fence: + // none (see this change's results). See memory + // `coopmat-lds-needs-explicit-memorybarriershared`. + memoryBarrierShared(); barrier(); [[unroll]] for (uint k = 0; k < WG_TILE_K / MMA_K; ++k) { @@ -403,6 +417,13 @@ void main() { [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { // Guards Csh against the previous iteration's readers. Inert on i == 0 but // must stay unconditional to remain workgroup-uniform. + // coopmat-lds-fence: barrier() alone does NOT order shared stores against a + // subsequent coopMatLoad on the M51 Xclipse/AMD-PAL driver -- symptom is one + // stale MMA_M-row band of A, all columns, ~2.5% of runs, no crash (observed + // 2026-09-02 in sdpa_compute_out_coopmat.glsl). Measured cost of the fence: + // none (see this change's results). See memory + // `coopmat-lds-needs-explicit-memorybarriershared`. + memoryBarrierShared(); barrier(); [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { #ifdef HAS_BIAS diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp index 3317b207fdf..09218f0fe9c 100644 --- a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp @@ -104,6 +104,11 @@ static const char* const kQ4gswTsweepPrefixes[] = { "tsweep_dbuf2_t", "tsweep_dbuf3_t", "tsweep_dbuf4_t", + // (coopmat-lds-fence 2026-09-03: tsweep_dbuf4nf_t / tsweep_dbuf4zpgtrnf_t + // -- unfenced clones of the two shipped defaults, used to measure the cost + // of the memoryBarrierShared() fix via an interleaved same-binary A/B. + // Result: 8da4w +0.001%, 4w -0.002%, both inside a 0.017-0.037% noise band. + // Deleted after measurement.) "tsweep_t", }; From 7f43322dc19428150cdc47860d089b31ecc759f6 Mon Sep 17 00:00:00 2001 From: Yanwen Xu Date: Wed, 2 Sep 2026 21:24:06 -0700 Subject: [PATCH 17/28] [ET-VK] Fix intermittent wrong output in SDPA attn*V coopmat sdpa_compute_out_coopmat produced silently wrong output on ~2.5% of runs at S=256, with no crash and no DEVICE_LOST. The previous correctness gate could not see it: both its shapes were S=128, where the shader runs only 2 M-tiles and 4 K-chunks, and the failure needs more. Diagnosis. The microbench seeds with srand(0), so every repeat feeds byte identical inputs to an identical shader on an identical driver -- and the outcome still varied run to run. That rules out a compile-time miscompile (those fail every time, as the QK^T UBO-stride bug did at 20/20) and means a race. Localizing the wrong elements pinned it: each failure corrupted exactly one MMA_M-row band -- 16 consecutive rows, ALL WG_TILE_N columns, a single head. Rows confined to one 16-row block while every column is wrong is the signature of a bad A operand, since A is indexed by row and shared across all column blocks. Those rows are Ash rows written by one subgroup and read back by coopMatLoad in several, and the band was stale even for the subgroup that wrote it -- so the store was late, not misindexed. barrier() alone was not ordering the uvec4 Ash/Bsh stores against the cooperative-matrix load path on this driver. Make it explicit with memoryBarrierShared(), as the linear coopmat kernel already does for its shared staging. Evidence: 0 failures / 300 case-runs with attn*V coopmat isolated (150 reps), and 0 / 160 on the full extended gate with both coopmat shaders (40 reps). Against the measured pre-fix rate of 5/200, zero failures in 300 has probability ~0.05%. Per-kernel cost is nil: attn*V prefill is 13002.16 us with the fence against 13002.20 us without (n=5, CoV 0.008%). Residual risk, stated rather than assumed away: only the write->read barrier got the fence. The loop's second barrier orders this chunk's coopMatLoad reads against the next chunk's stores, and if the same driver quirk applies in that direction it would reproduce an identical symptom. 460 clean case-runs do not exclude a rarer WAR variant at a ~2.5% base rate. (cherry picked from commit 8dc05ca8967ed5f0987e08e2839b63935d6010c9) --- .../graph/ops/glsl/sdpa_compute_out_coopmat.glsl | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coopmat.glsl b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coopmat.glsl index d3173441ccb..ed6f13f206f 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coopmat.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coopmat.glsl @@ -193,6 +193,17 @@ void main() { packFloat2x16(v1.xy), packFloat2x16(v1.zw)); } + // The Ash/Bsh stores above are read back by coopMatLoad below, and by + // OTHER subgroups than the one that wrote them (Ash row r is written by + // invocation r*INVS_PER_ROW_A.. and read by every subgroup whose + // warpInTile.y covers r). barrier() alone was not ordering those uvec4 + // stores against the cooperative-matrix load path on the Xclipse/AMD-PAL + // driver: ~2.5% of runs at S=256 produced exactly one stale MMA_M-row + // band of A -- 16 consecutive rows, all WG_TILE_N columns, one head -- + // which is a lost/late Ash write, not a wrong index (the write was + // invisible even to its own writer). Make the shared-memory ordering + // explicit, as the linear coopmat kernel does. + memoryBarrierShared(); barrier(); // --- Cooperative matrix MMA (identical to coopmat_mm) --- From 65dfa1c151e419efc8ef0e8d99e951028e76e463 Mon Sep 17 00:00:00 2001 From: Yanwen Xu Date: Wed, 2 Sep 2026 18:35:05 -0700 Subject: [PATCH 18/28] [ET-VK] SDPA coopmat correctness gate: reach every causal-mask region The SDPA coopmat gate had two cases, both S=128. At the shipped 128x64 QK^T tile that gives num_tiles_m=1, num_tiles_n=2, so BOTH tiles are diagonal: the shader's whole-tile all-masked early-out never fired and an all-visible fast path would have gotten zero coverage. A gate cannot be treated as covering a path its shapes cannot reach, so extend it before touching any shader. --sdpa-regions-only enumerates the QK^T tile grid and prints each tile's region classification, then reports which of the three shader paths the gate actually reaches. It also checks the classification at its boundaries: for a fixed M-tile row the classes must walk visible -> diagonal -> masked with no class recurring and no column skipped, which catches both a tile claimed by two classes and a tile claimed by none. Two S=256 cases are added as a separate "regions" tier. S=256 is the smallest size reaching all three regions at the shipped tile (num_tiles_m=2, num_tiles_n=4), and it populates both boundary transitions; S=192 is not usable because it is not a multiple of WG_TILE_M=128 and the dispatch gate would refuse the shader. It also gives attn*V 8 K-chunks instead of 4. --sdpa-tier= keeps the cheap S=128 pre-check available; all three tiers measure 222/355/406 ms per pass, so the 10+ repeat discipline stays practical even for the full gate. --sdpa-force-fallback runs the same shapes with coopmat disabled. This proves the gate's dispatch assertion is load-bearing (the tiled path is numerically correct, so it reports mismatches=0 and the gate must still fail it), and it is the control that separates a wrong shader from a wrong harness. Extending the gate immediately exposed a pre-existing intermittent failure in sdpa_compute_out_coopmat at S=256 (4/100 case-runs) with QK^T clean at 0/100, and the tiled path clean at 0/40 on identical shapes. The two ET_VK_SDPA_DISABLE_{QK,OUT}_COOPMAT switches here are the diagnostic that attributed it and are needed to run a QK^T-isolated gate while attn*V remains flaky. They are temporary scaffolding and must be removed before promotion. (cherry picked from commit 9d31d99d73022b106083769c9c0ca431fab38600) --- .../vulkan/runtime/graph/ops/impl/SDPA.cpp | 5 + .../test/custom_ops/test_llama_microbench.cpp | 320 +++++++++++++++++- 2 files changed, 320 insertions(+), 5 deletions(-) diff --git a/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp b/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp index 9fe48b6fdec..bb911a9d850 100644 --- a/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp @@ -349,7 +349,11 @@ vkapi::ShaderInfo pick_sdpa_qk_shader( const bool is_gemv = is_single_token(graph, q_projected); // Prefill WMMA path: Q @ K^T with K = head_dim, N = context_len, M = S. + // DEBUG BISECT (temporary): disable coopmat for ONE of the two SDPA GEMMs + // so an intermittent numerical failure can be attributed to a specific + // shader. Reverted once the attribution is recorded. if (!is_gemv && sdpa_coopmat_device_ok(graph) && + std::getenv("ET_VK_SDPA_DISABLE_QK_COOPMAT") == nullptr && sdpa_buf_half(graph, q_projected) && sdpa_buf_half(graph, k_cache) && sdpa_buf_half(graph, attn_weights)) { const SDPADims d = compute_sdpa_dims( @@ -503,6 +507,7 @@ vkapi::ShaderInfo pick_sdpa_av_shader( // Prefill WMMA path: P @ V with K = context_len, N = head_dim, M = S. if (!is_gemv && sdpa_coopmat_device_ok(graph) && + std::getenv("ET_VK_SDPA_DISABLE_OUT_COOPMAT") == nullptr && sdpa_buf_half(graph, out) && sdpa_buf_half(graph, attn_weights_softmax) && sdpa_buf_half(graph, v_cache)) { diff --git a/backends/vulkan/test/custom_ops/test_llama_microbench.cpp b/backends/vulkan/test/custom_ops/test_llama_microbench.cpp index 30ba08e33c1..2b548bc0d6e 100644 --- a/backends/vulkan/test/custom_ops/test_llama_microbench.cpp +++ b/backends/vulkan/test/custom_ops/test_llama_microbench.cpp @@ -34,6 +34,20 @@ // Other flags: // --model= only run models whose name contains // --correctness-only run just the linear correctness matrix, skip perf +// --sdpa-regions-only enumerate the QK^T causal-mask tile grid for every +// SDPA correctness case and print each tile's region +// classification (all_masked / all_visible / diagonal), +// then report which of the three shader paths the gate +// actually reaches. Host-side only, no GPU. A path no +// case produces a tile for is UNCOVERED, however many +// times --sdpa-correctness-only passes. +// --sdpa-tier= +// which SDPA correctness tier to run; default all. +// "fast" = the original S=128 cases, the cheap +// post-edit pre-check. "regions" = the S=256 cases +// that put a tile in every QK^T mask region (4x the +// reference cost). "all" = the extended gate, which is +// what an accept decision requires. // --sdpa-correctness-only run just the SDPA coopmat correctness cases // (sdpa_compute_attn_weights_coopmat / // sdpa_compute_out_coopmat vs. a CPU causal-attention @@ -1396,19 +1410,262 @@ struct SdpaCorrectnessCase { int64_t head_dim; // D int64_t num_heads; // Q_H int64_t num_kv_heads; // KV_H + // "fast" -- the cheap pre-check tier, run after any shader edit. + // "regions" -- the region-coverage tier: the smallest shapes that put a + // tile in EVERY QK^T mask region (see sdpa_report_qk_regions). + // 4x the reference cost of "fast", so it is kept separate to + // protect the 10+ repeat discipline (--sdpa-tier). + const char* tier; }; const std::vector kSdpaCorrectnessCases = { // Minimal GQA case: 128 is the smallest legal QK^T M-tile multiple, 64 // the smallest legal head_dim (both QK^T K-tile and attn*V N-tile). - {"tiny_gqa", 128, 64, 2, 1}, + {"tiny_gqa", 128, 64, 2, 1, "fast"}, // 1B's real head configuration (head_dim=64, 32 Q heads, 8 KV heads -- // kSdpaModels), S truncated from the real 2048 to the same aligned 128 // so the CPU reference stays fast; this is the shape most likely to // exercise a head-indexing (GQA) bug the tiny case's group size of 2 // could hide. - {"1b_head_config", 128, 64, 32, 8}, + {"1b_head_config", 128, 64, 32, 8, "fast"}, + // Region-coverage tier. S=256 is the SMALLEST size that reaches all three + // QK^T mask regions at the shipped 128x64 tile: it gives num_tiles_m=2, + // num_tiles_n=4, so + // m=0: n=0,1 diagonal n=2,3 fully masked + // m=1: n=0,1 fully visible n=2,3 diagonal + // which populates every class AND both boundary transitions + // (visible->diagonal and diagonal->masked). S=192 is not usable: it is + // not a multiple of WG_TILE_M=128, so SDPA.cpp's alignment gate would + // refuse to dispatch the coopmat shader at all. + // + // Also gives sdpa_out 8 K-chunks (context_len/WG_TILE_K = 256/32) instead + // of 4, so a double-buffer ping-pong parity error cannot hide in a + // 2-iteration loop. + // + // Cost: the reference is O(S^2*D*Q_H), so these are 4x the "fast" cases + // (~537 M MACs for the 32-head one against ~134 M) -- seconds per pass, + // not sub-second. That is why they are a separate tier. + {"tiny_gqa_s256", 256, 64, 2, 1, "regions"}, + {"1b_head_config_s256", 256, 64, 32, 8, "regions"}, }; +// ---------------- QK^T mask-region enumeration (host-side) ---------------- +// sdpa_compute_attn_weights_coopmat.glsl classifies each WG tile by where it +// sits relative to the causal diagonal, and runs different code per class. +// A correctness gate only covers a class if some case actually produces a +// tile in it, so enumerate the grid rather than assume (specs +// sdpa-coopmat-causal-mask-paths: "A correctness gate is not treated as +// covering a path its shapes cannot reach"). +// +// fully masked : c_tile_base > s_tile_base + WG_TILE_M - 1 + input_pos +// (lowest context index in the tile already exceeds the +// highest s + input_pos) -- the shader's tile_all_masked +// fully visible: c_tile_base + WG_TILE_N - 1 <= s_tile_base + input_pos +// (highest context index is within the lowest row's window) +// diagonal : neither, so the per-element mask is required +// +// Both fast-path conditions holding at once would require +// WG_TILE_M + WG_TILE_N < 2, so the classification is exhaustive and +// non-overlapping by construction -- `both` below asserts that per tile +// instead of trusting the algebra. +// +// Tile dims mirror SDPA.cpp's kSdpaAttnDefaultDims / the shipped default in +// sdpa_compute_attn_weights_coopmat.yaml. An ET_VK_SDPA_ATTN_COOPMAT_VARIANT +// override would need these updated to match. +constexpr int64_t kSdpaAttnWgTileM = 128; +constexpr int64_t kSdpaAttnWgTileN = 64; + +struct SdpaRegionCounts { + int64_t tiles = 0; + int64_t all_masked = 0; + int64_t all_visible = 0; + int64_t diagonal = 0; + int64_t both = 0; // overlap; must stay 0 +}; + +// Enumerates the QK^T tile grid for one shape. `verbose` prints the class of +// every tile (task 2.1 wants the per-tile classification, not just totals). +SdpaRegionCounts sdpa_enumerate_qk_regions( + int64_t S, + int64_t context_len, + int64_t input_pos, + bool verbose, + int64_t wg_tile_m = kSdpaAttnWgTileM, + int64_t wg_tile_n = kSdpaAttnWgTileN) { + SdpaRegionCounts c; + const int64_t num_tiles_m = (S + wg_tile_m - 1) / wg_tile_m; + const int64_t num_tiles_n = (context_len + wg_tile_n - 1) / wg_tile_n; + for (int64_t i = 0; i < num_tiles_m; ++i) { + for (int64_t j = 0; j < num_tiles_n; ++j) { + const int64_t s_base = wg_tile_m * i; + const int64_t c_base = wg_tile_n * j; + const bool masked = c_base > s_base + wg_tile_m - 1 + input_pos; + const bool visible = c_base + wg_tile_n - 1 <= s_base + input_pos; + ++c.tiles; + const char* label; + if (masked && visible) { + ++c.both; + label = "BOTH(BUG)"; + } else if (masked) { + ++c.all_masked; + label = "all_masked"; + } else if (visible) { + ++c.all_visible; + label = "all_visible"; + } else { + ++c.diagonal; + label = "diagonal"; + } + if (verbose) { + std::cout << "[sdpa-regions] tile m=" << i << " n=" << j + << " s_base=" << s_base << " c_base=" << c_base << " -> " + << label << "\n"; + } + } + } + return c; +} + +// Checks the classification is exhaustive and non-overlapping AT ITS +// BOUNDARIES, which is where an off-by-one actually hides. For a fixed M-tile +// row, increasing the N-tile column must walk the classes in exactly the order +// [all_visible...] [diagonal...] [all_masked...] +// with no class recurring once left. That single property catches both failure +// modes the spec names: a boundary tile claimed by two classes (`overlap`, +// which the shader's two conditions would have to both accept) and a boundary +// tile claimed by none (which would show up as a class reappearing after the +// run it belongs to). The transition columns themselves are printed, so the +// "last fully-visible / first diagonal" and "last diagonal / first +// fully-masked" pairs are on the record rather than inferred. +bool sdpa_check_region_boundaries( + int64_t S, + int64_t context_len, + int64_t input_pos, + const char* label) { + const int64_t num_tiles_m = (S + kSdpaAttnWgTileM - 1) / kSdpaAttnWgTileM; + const int64_t num_tiles_n = + (context_len + kSdpaAttnWgTileN - 1) / kSdpaAttnWgTileN; + bool ok = true; + for (int64_t i = 0; i < num_tiles_m; ++i) { + // phase 0 = expecting all_visible, 1 = diagonal, 2 = all_masked + int phase = 0; + int64_t last_visible = -1, first_diagonal = -1; + int64_t last_diagonal = -1, first_masked = -1; + for (int64_t j = 0; j < num_tiles_n; ++j) { + const int64_t s_base = kSdpaAttnWgTileM * i; + const int64_t c_base = kSdpaAttnWgTileN * j; + const bool masked = c_base > s_base + kSdpaAttnWgTileM - 1 + input_pos; + const bool visible = c_base + kSdpaAttnWgTileN - 1 <= s_base + input_pos; + if (masked && visible) { + std::cout << "[sdpa-boundary] " << label << " m=" << i << " n=" << j + << " OVERLAP: satisfies BOTH all_masked and all_visible\n"; + ok = false; + continue; + } + const int cls = masked ? 2 : (visible ? 0 : 1); + if (cls < phase) { + std::cout << "[sdpa-boundary] " << label << " m=" << i << " n=" << j + << " OUT OF ORDER: class " << cls + << " reappeared after phase " << phase + << " -- classification is not a clean visible/diagonal/" + "masked partition\n"; + ok = false; + } + phase = cls > phase ? cls : phase; + if (cls == 0) { + last_visible = j; + } else if (cls == 1) { + if (first_diagonal < 0) { + first_diagonal = j; + } + last_diagonal = j; + } else if (first_masked < 0) { + first_masked = j; + } + } + std::cout << "[sdpa-boundary] " << label << " m=" << i + << " last_visible=" << last_visible + << " first_diagonal=" << first_diagonal + << " last_diagonal=" << last_diagonal + << " first_masked=" << first_masked; + // Adjacency: a present transition must be between consecutive columns, + // i.e. no column is skipped between one class's end and the next's start. + bool adjacent = true; + if (last_visible >= 0 && first_diagonal >= 0 && + first_diagonal != last_visible + 1) { + adjacent = false; + } + if (last_diagonal >= 0 && first_masked >= 0 && + first_masked != last_diagonal + 1) { + adjacent = false; + } + if (!adjacent) { + std::cout << " NON-ADJACENT TRANSITION (a column is unclassified)"; + ok = false; + } + std::cout << "\n"; + } + return ok; +} + +// Prints the region distribution for every correctness case and reports +// whether the gate as a whole reaches all three classes. Returns true iff +// no tile is doubly-classified AND all three classes are covered. +bool sdpa_report_qk_regions(bool verbose, const char* tier = "all") { + SdpaRegionCounts total; + bool boundaries_ok = true; + for (const auto& c : kSdpaCorrectnessCases) { + if (std::string(tier) != "all" && std::string(c.tier) != tier) { + continue; + } + // input_pos == 0 for every case, so context_len == seq_len. + const SdpaRegionCounts r = + sdpa_enumerate_qk_regions(c.seq_len, c.seq_len, 0, verbose); + std::cout << "[sdpa-regions] " << c.name << " S=" << c.seq_len + << " context_len=" << c.seq_len << " tile=" << kSdpaAttnWgTileM + << "x" << kSdpaAttnWgTileN << " num_tiles_m=" + << (c.seq_len + kSdpaAttnWgTileM - 1) / kSdpaAttnWgTileM + << " num_tiles_n=" + << (c.seq_len + kSdpaAttnWgTileN - 1) / kSdpaAttnWgTileN + << " tiles=" << r.tiles << " all_masked=" << r.all_masked + << " all_visible=" << r.all_visible << " diagonal=" << r.diagonal + << " overlap=" << r.both << "\n"; + total.tiles += r.tiles; + total.all_masked += r.all_masked; + total.all_visible += r.all_visible; + total.diagonal += r.diagonal; + total.both += r.both; + boundaries_ok = + sdpa_check_region_boundaries(c.seq_len, c.seq_len, 0, c.name) && + boundaries_ok; + } + const bool exhaustive = + total.all_masked + total.all_visible + total.diagonal == total.tiles; + const bool covered = + total.all_masked > 0 && total.all_visible > 0 && total.diagonal > 0; + std::cout << "[sdpa-regions] TOTAL tiles=" << total.tiles + << " all_masked=" << total.all_masked + << " all_visible=" << total.all_visible + << " diagonal=" << total.diagonal << " overlap=" << total.both + << " exhaustive=" << (exhaustive ? "yes" : "NO") + << " all_three_covered=" << (covered ? "yes" : "NO") << "\n"; + if (!covered) { + std::cout << "[sdpa-regions] UNCOVERED:"; + if (total.all_masked == 0) { + std::cout << " all_masked"; + } + if (total.all_visible == 0) { + std::cout << " all_visible"; + } + if (total.diagonal == 0) { + std::cout << " diagonal"; + } + std::cout << " -- these shader paths are NOT covered by this gate\n"; + } + std::cout << "[sdpa-regions] boundaries_ok=" << (boundaries_ok ? "yes" : "NO") + << "\n"; + return total.both == 0 && exhaustive && covered && boundaries_ok; +} + // Causal, GQA-aware fp32 CPU reference. q is [S, Q_H, D], k/v are // [S, KV_H, D] (row-major, batch=1 squeezed). kv_h = q_h / (Q_H / KV_H), // matching sdpa_compute_attn_weights_coopmat.glsl's GQA head mapping exactly @@ -1463,8 +1720,21 @@ std::vector sdpa_reference( // silent tiled fallback -- tiled is also numerically correct, so a pure // value comparison alone cannot tell the two apart) AND every output // element is within tolerance. +// --sdpa-force-fallback: run the SDPA correctness cases with coopmat DISABLED, +// so the tiled shaders serve the same shapes. Two uses: +// 1. it proves the gate's dispatch assertion is load-bearing -- the tiled path +// is also numerically correct, so a pass here with qk_coopmat=NO must still +// be reported FAILED, otherwise a silent fallback would look like a pass; +// 2. it is the control that separates "the coopmat shader is wrong" from "the +// harness/reference/softmax is wrong" when a case fails intermittently. +bool g_sdpa_force_fallback = false; + bool sdpa_correctness_case(const SdpaCorrectnessCase& c) { - unsetenv("ET_VK_DISABLE_COOPMAT"); + if (g_sdpa_force_fallback) { + setenv("ET_VK_DISABLE_COOPMAT", "1", 1); + } else { + unsetenv("ET_VK_DISABLE_COOPMAT"); + } GraphConfig config; config.enable_querypool = true; @@ -1598,6 +1868,7 @@ bool sdpa_correctness_case(const SdpaCorrectnessCase& c) { << " ref=" << ref[first_mismatch] << ")"; } std::cout << (numeric_ok && fired_ok ? " PASSED" : " FAILED") << "\n"; + unsetenv("ET_VK_DISABLE_COOPMAT"); // restore the tree's default-on state return numeric_ok && fired_ok; } @@ -1607,11 +1878,32 @@ bool sdpa_correctness_case(const SdpaCorrectnessCase& c) { // multiple times in a loop -- kept a single pass per call, like // run_linear_correctness, so a driver script controls the rep count and can // distinguish "which specific repeat failed." -bool run_sdpa_correctness() { +bool run_sdpa_correctness(const char* tier = "all") { bool all_ok = true; + // Report which QK^T mask regions these shapes actually reach before running + // them, so a pass is never mistaken for coverage of a path no case produces + // a tile for. Non-verbose: --sdpa-regions-only prints the per-tile detail. + const bool regions_covered = sdpa_report_qk_regions(/*verbose=*/false, tier); + if (!regions_covered) { + std::cout << "[sdpa-correctness] WARNING: tier=" << tier + << " does not cover every QK^T mask-region path (see " + "[sdpa-regions] above)\n"; + } + int64_t ran = 0; for (const auto& c : kSdpaCorrectnessCases) { + if (std::string(tier) != "all" && std::string(c.tier) != tier) { + continue; + } + ++ran; all_ok = sdpa_correctness_case(c) && all_ok; } + std::cout << "[sdpa-correctness] tier=" << tier << " cases_run=" << ran + << "\n"; + if (ran == 0) { + std::cout << "[sdpa-correctness] no case matched tier '" << tier + << "' -- treating as failure rather than a silent pass\n"; + return false; + } return all_ok; } @@ -1968,6 +2260,12 @@ void print_usage() { " --correctness-only run just the linear correctness matrix\n" " --sdpa-correctness-only run just the SDPA coopmat correctness " "cases\n" + " --sdpa-regions-only enumerate the QK^T mask-region tile grid " + "(no GPU)\n" + " --sdpa-tier= which SDPA correctness tier to " + "run (default all)\n" + " --sdpa-force-fallback run SDPA correctness with coopmat " + "DISABLED (control)\n" " --skip-correctness skip the correctness gate before perf\n" " --list print every case with its sizes, no GPU\n" " --help this message\n"; @@ -2024,6 +2322,8 @@ void list_cases( int main(int argc, char** argv) { bool linear = false, baseline = false, sdpa = false; bool correctness_only = false, sdpa_correctness_only = false; + bool sdpa_regions_only = false; + std::string sdpa_tier = "all"; bool skip_correctness = false, list_only = false; bool production_diff = false; // Additive machine-readable output. Absent, every existing line is byte @@ -2043,6 +2343,12 @@ int main(int argc, char** argv) { correctness_only = true; } else if (arg == "--sdpa-correctness-only") { sdpa_correctness_only = true; + } else if (arg == "--sdpa-regions-only") { + sdpa_regions_only = true; + } else if (arg.rfind("--sdpa-tier=", 0) == 0) { + sdpa_tier = arg.substr(std::string("--sdpa-tier=").size()); + } else if (arg == "--sdpa-force-fallback") { + g_sdpa_force_fallback = true; } else if (arg == "--production-diff") { production_diff = true; } else if (arg == "--skip-correctness") { @@ -2142,8 +2448,12 @@ int main(int argc, char** argv) { if (correctness_only) { return finish_correctness(run_linear_correctness(filter)); } + if (sdpa_regions_only) { + return finish_correctness( + sdpa_report_qk_regions(/*verbose=*/true, sdpa_tier.c_str())); + } if (sdpa_correctness_only) { - return finish_correctness(run_sdpa_correctness()); + return finish_correctness(run_sdpa_correctness(sdpa_tier.c_str())); } if (production_diff) { return finish_correctness(run_production_diff()); From df5b8dca1beb2ee4eba22059b040e9737e05cf87 Mon Sep 17 00:00:00 2001 From: Yanwen Xu Date: Wed, 2 Sep 2026 18:35:42 -0700 Subject: [PATCH 19/28] [ET-VK] SDPA QK^T coopmat: all-visible whole-tile fast path At the 2048-prefill workload with the shipped 128x64 tile there are 512 QK^T workgroups: 240 are entirely above the diagonal and already early-out, 32 genuinely straddle it, and 240 are entirely BELOW it -- fully visible, nothing to mask. Those 240 were still paying the full masked epilogue: a coopMatStore into the Csh scratch, a barrier, and a per-element scalar read-mask-write over the whole 128x64 tile. That is 240 of the 272 workgroups that execute (88%) doing an LDS round trip that cannot mask anything. Classify the tile and send the fully-visible case straight to global, as the linear kernel's buffer epilogue does. The fully-visible condition is the complement of the existing all-masked one; both holding at once would require WG_TILE_M + WG_TILE_N < 2, and a tile matching neither falls through to the per-element path, so the three classes are exhaustive and disjoint. The address arithmetic reproduces the Csh path's global address exactly. The store's stride cannot be the UBO-derived aw_row_width. coopMatStore miscompiles on the Xclipse/AMD-PAL compiler when its stride derives from a UBO value -- the same bug sdpa_compute_out_coopmat.glsl works around for its output stride -- and measured deterministically here: 20/20 case-runs wrong with the UBO value, 0/20 with a compile-time one. Unlike the linear kernel's out_N_arg this width is not static; resize_sdpa_attn_weights_node recomputes it from the input_pos symint on every resize. So it is passed as a spec constant resolved at node construction and the fast path is entered ONLY when that baked value still equals the live width, falling through to the compile-time-stride Csh path otherwise. Correctness therefore never depends on the guess: a chunked prefill simply does not get the fast path. The entry condition also requires the tile to lie wholly inside both extents, since the store writes whole MMA tiles unchecked. The dispatch gate already guarantees that (and this shader's staging reads depend on it unguarded), so under that gate no workgroup loses the fast path -- but it makes an out-of-extent write impossible by construction rather than by appeal to the gate. Per-kernel QK^T prefill time, 3 suite reps each, dispatch confirmed every rep: 8B 9505.6 -> 8178.6 us (-14.0%), 3B 7199.2 -> 6194.0 us (-14.0%), 1B 6806.5 -> 6167.4 us (-9.4%). Between-rep spread is 0.05-0.30% for 8B/3B, two orders of magnitude below the effect; 1B's fast-path reps spread 6.4% (5922.8-6319.8) so its magnitude is not resolved by 3 reps, only its direction. Correctness: 0/60 case-runs over 15 reps of the extended gate with attn*V held on the tiled path (its coopmat shader has a pre-existing intermittent failure at S=256, unrelated to this change). The fast path is confirmed live rather than assumed: perturbing only its store fails S=256 and leaves S=128 untouched -- matching the enumerated claim that no S=128 tile is fully visible -- with both first mismatches at s=128, the first row of the fully-visible M-tile row. (cherry picked from commit 013ef5a75d6006311229876be96d441d52875b6f) --- .../sdpa_compute_attn_weights_coopmat.glsl | 91 +++++++++++++++++++ .../vulkan/runtime/graph/ops/impl/SDPA.cpp | 31 ++++++- 2 files changed, 119 insertions(+), 3 deletions(-) diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_coopmat.glsl b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_coopmat.glsl index ac24c1d8d59..fe95e6cb605 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_coopmat.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_coopmat.glsl @@ -31,6 +31,37 @@ * per-element mask. A whole-WG-tile that is entirely above the diagonal is * written as -inf and skips the MMA loop (~halves prefill QK^T work). * + * Each output tile is classified by where it sits relative to the causal + * diagonal, and only the diagonal class needs the Csh round trip: + * all masked : c_tile_base > s_tile_base + WG_TILE_M - 1 + input_pos + * all visible : c_tile_base + WG_TILE_N - 1 <= s_tile_base + input_pos + * diagonal : neither + * Both whole-tile conditions holding at once would need + * WG_TILE_M + WG_TILE_N < 2, and a tile matching neither falls through to the + * per-element path, so the classes are exhaustive and disjoint. + * + * Region counts at the reported 2048-prefill workload (input_pos = 0, 128x64 + * tile, so 16 M-tiles x 32 N-tiles = 512 workgroups): + * 240 all masked (early-out, no MMA) + * 240 all visible (direct coopMatStore to global: no Csh, no barrier, no + * per-element mask) + * 32 diagonal (per-element mask via Csh) + * 240 + 240 + 32 = 512, so the all-visible path covers 240 of the 272 + * workgroups that actually execute (88%). + * + * WHY THE ALL-VISIBLE STORE TAKES ITS STRIDE FROM A SPEC CONSTANT: + * coopMatStore's stride operand must not derive from a UBO value on the + * Xclipse/AMD-PAL compiler -- the same miscompile sdpa_compute_out_coopmat.glsl + * works around for its output stride. attn_weights' row width is + * align_up_4(context_len), and context_len = input_pos + S comes from UBOs, so + * using it directly produces silently wrong output (measured: deterministic, + * ~3500-5800 wrong elements per case; with a compile-time stride, zero). + * Unlike the linear kernel's out_N_arg, this width is NOT static -- it is + * recomputed on every resize (SDPA.cpp resize_sdpa_attn_weights_node) -- so the + * spec constant is a BAKED GUESS and the fast path is taken only when it + * matches the live width. Otherwise control falls through to the Csh path, + * whose stride is the compile-time WG_TILE_N and is therefore always safe. + * * Dispatch: global {num_tiles_n*WG_SIZE, num_tiles_m, H_q}, local {WG_SIZE,1,1}. * tileID = gl_WorkGroupID.xy (x->context, y->seq), q_h = gl_WorkGroupID.z. */ @@ -70,6 +101,10 @@ ${layout_declare_spec_const(C, "float", "inv_scale", "1.0")} // Xclipse/AMD-PAL compiler crashes on a coopMatMulAdd loop with a UBO-derived // trip count — see coopmat_mm.glsl). ${layout_declare_spec_const(C, "int", "num_k_chunks_arg", "0")} +// attn_weights row width (= align_up_4(context_len)) as resolved at node +// construction. Only the all-visible fast path reads it, and only after +// confirming it equals the live UBO-derived width -- see the file header. +${layout_declare_spec_const(C, "int", "aw_row_width_arg", "0")} const uint MMA_M = ${MMA_M}; const uint MMA_N = ${MMA_N}; @@ -145,6 +180,28 @@ void main() { // the highest (s + input_pos), every element is masked. const bool tile_all_masked = int(c_tile_base) > (int(s_tile_base) + int(WG_TILE_M) - 1 + input_pos); + + // Whole-tile fully-visible fast path (complement of tile_all_masked; see + // the file header for exhaustiveness and the region counts). Three + // conditions, all workgroup-uniform: + // 1. nothing in the tile is masked; + // 2. the tile lies wholly inside both extents -- the store writes whole + // MMA tiles with no per-element bound check. SDPA.cpp's dispatch gate + // already guarantees this (S % WG_TILE_M == 0 and + // context_len % WG_TILE_N == 0, which this shader's unguarded staging + // reads below also depend on), but stating it makes an out-of-extent + // write impossible by construction rather than by appeal to the gate, + // and costs no workgroup the fast path under that gate; + // 3. the baked stride matches the live row width, without which + // coopMatStore would need a UBO-derived stride -- see the file header. + const bool tile_in_extent = + (s_tile_base + WG_TILE_M <= uint(S)) && + (c_tile_base + WG_TILE_N <= uint(context_len)); + const bool aw_stride_is_static = aw_row_width_arg == aw_row_width; + const bool tile_all_visible = + tile_in_extent && aw_stride_is_static && + (int(c_tile_base) + int(WG_TILE_N) - 1 <= int(s_tile_base) + input_pos); + if (tile_all_masked) { for (uint idx = gl_LocalInvocationID.x; idx < WG_TILE_M * WG_TILE_N; idx += WG_SIZE) { @@ -240,6 +297,40 @@ void main() { barrier(); } + if (tile_all_visible) { + // Straight to global, mirroring the linear kernel's buffer epilogue + // (linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr.glsl:727-748). + // + // The address is identical to the Csh path's for the same element. + // Csh path: element (r, c) of MMA tile (i, j) lands in + // Csh[(local_row + r) * WG_TILE_N + local_col + c], then is copied to + // (q_h*S_aligned + s_tile_base + local_row + r) * aw_row_width + // + c_tile_base + local_col + c. + // Here: aw_tile_base + local_row*STRIDE + local_col + r*STRIDE + c, + // with aw_tile_base = (q_h*S_aligned + s_tile_base)*aw_row_width + // + c_tile_base, which expands to the same thing + // because STRIDE == aw_row_width is exactly what aw_stride_is_static + // checked. + const uint STRIDE = uint(aw_row_width_arg); + const uint aw_tile_base = + (uint(q_h) * uint(S_aligned) + s_tile_base) * STRIDE + c_tile_base; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + result[i][j] = result[i][j] * inv_scale; // fp32 scalar multiply + coopmat out_tile = + coopmat(result[i][j]); + const uint local_row = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + const uint local_col = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopMatStore( + out_tile, t_attn_weights, + aw_tile_base + local_row * STRIDE + local_col, + STRIDE, + gl_CooperativeMatrixLayoutRowMajor); + } + } + return; + } + // --- Scale on the fp32 accumulator, store fp16 into Csh [s][c] scratch --- const float16_t inv_scale_h = float16_t(inv_scale); [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { diff --git a/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp b/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp index bb911a9d850..2d2fde2f231 100644 --- a/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp @@ -672,6 +672,18 @@ void add_sdpa_compute_attn_weights_node( const ValueRef mode_ref = static_cast(mode); + // See the aw_row_width spec-constant note below. Mirrors + // resize_sdpa_attn_weights_node's own derivation so the two cannot disagree + // about what the row width is when they do agree on context_len. + int32_t aw_row_width_at_construction = 0; + if (mode == SDPAMode::LLM) { + const int32_t seq_len = graph.size_at(-3, q); + const int32_t input_pos_val = + is_valid(input_pos_symint) ? graph.read_symint(input_pos_symint) : 0; + aw_row_width_at_construction = + static_cast(utils::align_up_4(seq_len + input_pos_val)); + } + graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, pick_sdpa_qk_shader, @@ -683,15 +695,28 @@ void add_sdpa_compute_attn_weights_node( param_ubos, // Push Constants {}, - // Specialization Constants: {inv_scale (id 3), num_k_chunks (id 4)}. + // Specialization Constants: {inv_scale (id 3), num_k_chunks (id 4), + // aw_row_width (id 5)}. // num_k_chunks = head_dim / WG_TILE_K is static and consumed only by the // coopmat QK^T variant (WG_TILE_K from the active // ET_VK_SDPA_ATTN_COOPMAT_VARIANT); the tiled/coop variants declare - // only id 3 and ignore the trailing entry -- safe to compute + // only id 3 and ignore the trailing entries -- safe to compute // unconditionally even when coopmat doesn't end up firing. + // + // aw_row_width is attn_weights' row width, align_up_4(context_len), + // needed by the coopmat variant's all-visible fast path because + // coopMatStore's stride must not derive from a UBO value on the + // Xclipse/AMD-PAL compiler. Unlike the two above it is NOT static -- + // resize_sdpa_attn_weights_node recomputes it from the input_pos symint + // on every resize -- so this is the value at construction only, and the + // shader takes the fast path solely when it still matches the live + // width, falling back to its compile-time-stride Csh path otherwise. + // Correct for a single-shot prefill (input_pos == 0); a chunked prefill + // simply does not get the fast path. {scale_val, graph.size_at(-1, q) / - static_cast(sdpa_attn_tile_dims().k)}, + static_cast(sdpa_attn_tile_dims().k), + aw_row_width_at_construction}, // Resize Args: [q, k, input_pos_symint_or_dummy, mode] {q, k, input_pos_symint, mode_ref}, // Resizing Logic From 582ec23f95befafc05fcad2067c074c127e7f9da Mon Sep 17 00:00:00 2001 From: Yanwen Xu Date: Wed, 2 Sep 2026 22:23:08 -0700 Subject: [PATCH 20/28] [ET-VK] SDPA softmax: stop reducing over the causally-masked tail The row-wise softmax between the two SDPA GEMMs read and reduced over the whole context_len row, three times, for every row. But the QK^T shader has already written -inf to every element with c > s + input_pos, and a masked element cannot contribute: exp(-inf - max) is 0 in the sum, it cannot be the row max unless the entire row is masked, and it normalizes to exactly 0. So row s only needs the first (s + input_pos + 1) columns. Over a full prefill the rows are a triangle, not a rectangle, so this halves the bytes read. The store still covers the whole row, because attn*V stages every chunk with chunkK < context_len and multiplies it by V -- a stale tail would be multiplied as if it were attention weight. But the tail is written as zero WITHOUT loading the input, so pass 3 loses its tail read too, and the zero store is spread across all 64 workers rather than serialized on worker 0 the way the existing straddling-texel path is. Truncation is gated on HAS_INPUT_POS. Fused SDPA has no input_pos and takes its mask from an attn_mask bias instead, so no truncation is valid there. 8B prefill, per-kernel, n=5: 13971.1 -> 9451.9 us, -32.3%, with the two ranges disjoint ([13895.7,14048.7] vs [9405.0,9499.1]). Correctness 0 failures / 80 case-runs on the extended gate. Also adds a softmax bucket to test_llama_microbench's --sdpa suite, which only timed qk and av. That measurement is why this change is trustworthy: an ETDump put softmax at 443.3 -> 444.5 ms, i.e. no effect, but a single ETDump capture turned out to carry +-20% capture-to-capture noise (the same pair of captures moved attn*V by +20.5% across a change that the querypool measures at 0.00% with CoV 0.008%). The querypool path with repeats resolves a 32% effect the ETDump could not see at all. (cherry picked from commit 67facb0a41f801db7c429807d5b6e756dea9585b) --- .../ops/glsl/sdpa_attn_weights_softmax.glsl | 56 ++++++++++++++--- .../test/custom_ops/test_llama_microbench.cpp | 62 ++++++++++++++++++- 2 files changed, 106 insertions(+), 12 deletions(-) diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_attn_weights_softmax.glsl b/backends/vulkan/runtime/graph/ops/glsl/sdpa_attn_weights_softmax.glsl index 6c095e66255..3fcdaddd5cb 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/sdpa_attn_weights_softmax.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_attn_weights_softmax.glsl @@ -137,6 +137,30 @@ void main() { const int context_len_aligned_down = context_len - mod_4(context_len); const int C4_limit = div_4(context_len_aligned_down); + // Causal-mask truncation. In LLM mode the QK^T shader has already written + // -inf to every element with c > s + input_pos, so for row s only the first + // (s + input_pos + 1) columns carry information: a masked element + // contributes exp(-inf - max) == 0 to the sum, cannot be the row max unless + // the whole row is masked, and normalizes to exactly 0. All three passes can + // therefore stop at reduce_len instead of context_len, which on a full + // prefill halves the bytes read (the rows are a triangle, not a rectangle). + // + // The STORE still has to cover the whole context_len row: attn*V stages + // every chunk with chunkK < context_len and multiplies it by V, so the + // masked tail must be written as zero rather than left stale. Writing zero + // is cheaper than reading + exp + writing, and needs no load. + // + // Fused mode has no input_pos and gets its mask from an attn_mask bias + // instead, so no truncation is valid there. +#ifdef HAS_INPUT_POS + const int reduce_len = min(context_len, s + input_pos + 1); +#else + const int reduce_len = context_len; +#endif + const int reduce_texel_len = div_up_4(reduce_len); + const int reduce_len_aligned_down = reduce_len - mod_4(reduce_len); + const int R4_limit = div_4(reduce_len_aligned_down); + // ========================================================================= // Pass 1: Find the maximum value across the row for numerical stability. // Without this, exp(x) can overflow float32 when x > ~88.7. @@ -144,7 +168,7 @@ void main() { SOFTMAX_ACC_T local_max = SOFTMAX_ACC_T(-1.0 / 0.0); // -infinity - for (int c4 = worker_id; c4 < C4_limit; c4 += NUM_WORKERS_PER_WG) { + for (int c4 = worker_id; c4 < R4_limit; c4 += NUM_WORKERS_PER_WG) { SOFTMAX_IN_VEC4_T in_texel = load_attn_weights_c4( c4, s, q_h, context_texel_len, attn_S, Q_H); @@ -153,13 +177,13 @@ void main() { } } if (worker_id == 0) { - for (int c4 = C4_limit; c4 < context_texel_len; ++c4) { + for (int c4 = R4_limit; c4 < reduce_texel_len; ++c4) { const int c_base = mul_4(c4); SOFTMAX_IN_VEC4_T in_texel = load_attn_weights_c4( c4, s, q_h, context_texel_len, attn_S, Q_H); [[unroll]] for (int comp = 0; comp < 4; comp++) { - if (c_base + comp < context_len) { + if (c_base + comp < reduce_len) { local_max = max(local_max, SOFTMAX_ACC_T(in_texel[comp])); } } @@ -189,7 +213,7 @@ void main() { SOFTMAX_ACC_T local_exp_sum = SOFTMAX_ACC_T(0); - for (int c4 = worker_id; c4 < C4_limit; c4 += NUM_WORKERS_PER_WG) { + for (int c4 = worker_id; c4 < R4_limit; c4 += NUM_WORKERS_PER_WG) { SOFTMAX_IN_VEC4_T in_texel = load_attn_weights_c4( c4, s, q_h, context_texel_len, attn_S, Q_H); @@ -198,13 +222,13 @@ void main() { } } if (worker_id == 0) { - for (int c4 = C4_limit; c4 < context_texel_len; ++c4) { + for (int c4 = R4_limit; c4 < reduce_texel_len; ++c4) { const int c_base = mul_4(c4); SOFTMAX_IN_VEC4_T in_texel = load_attn_weights_c4( c4, s, q_h, context_texel_len, attn_S, Q_H); [[unroll]] for (int comp = 0; comp < 4; comp++) { - if (c_base + comp < context_len) { + if (c_base + comp < reduce_len) { local_exp_sum += exp(SOFTMAX_ACC_T(in_texel[comp]) - global_max); } } @@ -232,7 +256,8 @@ void main() { // Pass 3: Normalize each element: out = exp(x - max) / sum(exp(x - max)) // ========================================================================= - for (int c4 = worker_id; c4 < C4_limit; c4 += NUM_WORKERS_PER_WG) { + // Fully-inside-the-prefix texels: load, normalize, store. + for (int c4 = worker_id; c4 < R4_limit; c4 += NUM_WORKERS_PER_WG) { SOFTMAX_IN_VEC4_T in_texel = load_attn_weights_c4( c4, s, q_h, context_texel_len, attn_S, Q_H); @@ -244,15 +269,18 @@ void main() { store_attn_weights_softmax_c4( out_texel, c4, s, q_h, context_texel_len, attn_S, Q_H); } + + // The single texel that straddles reduce_len, if reduce_len is not a + // multiple of 4. Its masked lanes normalize to 0, same as the tail below. if (worker_id == 0) { - for (int c4 = C4_limit; c4 < context_texel_len; ++c4) { + for (int c4 = R4_limit; c4 < reduce_texel_len; ++c4) { const int c_base = mul_4(c4); SOFTMAX_IN_VEC4_T in_texel = load_attn_weights_c4( c4, s, q_h, context_texel_len, attn_S, Q_H); VEC4_T out_texel = VEC4_T(0); [[unroll]] for (int comp = 0; comp < 4; comp++) { - if (c_base + comp < context_len) { + if (c_base + comp < reduce_len) { out_texel[comp] = T( exp(SOFTMAX_ACC_T(in_texel[comp]) - global_max) / local_exp_sum); } @@ -261,4 +289,14 @@ void main() { out_texel, c4, s, q_h, context_texel_len, attn_S, Q_H); } } + + // Causally-masked tail: every element normalizes to exactly 0, so write zero + // WITHOUT loading the input. This is the read the truncation saves in pass 3, + // and the store attn*V depends on (it stages the whole context_len row). + // Spread across all workers, unlike the straddling texel above. + for (int c4 = reduce_texel_len + worker_id; c4 < context_texel_len; + c4 += NUM_WORKERS_PER_WG) { + store_attn_weights_softmax_c4( + VEC4_T(0), c4, s, q_h, context_texel_len, attn_S, Q_H); + } } diff --git a/backends/vulkan/test/custom_ops/test_llama_microbench.cpp b/backends/vulkan/test/custom_ops/test_llama_microbench.cpp index 2b548bc0d6e..0c64bf669d2 100644 --- a/backends/vulkan/test/custom_ops/test_llama_microbench.cpp +++ b/backends/vulkan/test/custom_ops/test_llama_microbench.cpp @@ -96,6 +96,7 @@ #include #include #include +#include #include #include #include @@ -1122,12 +1123,19 @@ const std::vector kSdpaRegimes = { }; struct SdpaRunResult { - float mean_us; // total (qk + av) + float mean_us; // total (qk + softmax + av) float stdev_us; float qk_mean_us; float qk_stdev_us; float av_mean_us; float av_stdev_us; + // The row-wise softmax between the two GEMMs. Timed separately because an + // ETDump showed it is the LARGEST single SDPA shader at the 8B prefill shape + // -- bigger than either GEMM -- and a single ETDump capture is far too noisy + // to optimize against (capture-to-capture spread was seen at +-20%, against + // a querypool CoV here of well under 1%). + float softmax_mean_us; + float softmax_stdev_us; std::vector dispatched_kernels; // from the last timed run }; @@ -1263,6 +1271,7 @@ SdpaRunResult sdpa_run_case( std::vector total_timings_us; std::vector qk_timings_us; std::vector av_timings_us; + std::vector softmax_timings_us; std::vector last_dispatched; for (int i = 0; i < kTimedRuns; ++i) { graph.execute(); @@ -1272,11 +1281,19 @@ SdpaRunResult sdpa_run_case( float qk_time_us = 0.0f; float av_time_us = 0.0f; + float softmax_time_us = 0.0f; last_dispatched.clear(); for (const auto& r : shader_results) { last_dispatched.push_back(r.kernel_name); const uint64_t duration_ns = r.end_time_ns - r.start_time_ns; - if (r.kernel_name.find("sdpa_compute_attn_weights") != + // Order matters: "sdpa_attn_weights_softmax" also contains + // "sdpa_attn_weights", but NOT "sdpa_compute_attn_weights", so the qk + // test below cannot capture it. Checked first regardless, so a future + // rename cannot silently fold softmax into the qk bucket. + if (r.kernel_name.find("softmax") != std::string::npos) { + softmax_time_us += static_cast(duration_ns) / 1000.0f; + } else if ( + r.kernel_name.find("sdpa_compute_attn_weights") != std::string::npos) { qk_time_us += static_cast(duration_ns) / 1000.0f; } else if (r.kernel_name.find("sdpa_compute_out") != std::string::npos) { @@ -1285,7 +1302,8 @@ SdpaRunResult sdpa_run_case( } qk_timings_us.push_back(qk_time_us); av_timings_us.push_back(av_time_us); - total_timings_us.push_back(qk_time_us + av_time_us); + softmax_timings_us.push_back(softmax_time_us); + total_timings_us.push_back(qk_time_us + av_time_us + softmax_time_us); } SdpaRunResult result; @@ -1295,6 +1313,9 @@ SdpaRunResult sdpa_run_case( result.qk_stdev_us = stdev_of(qk_timings_us, result.qk_mean_us); result.av_mean_us = mean_of(av_timings_us); result.av_stdev_us = stdev_of(av_timings_us, result.av_mean_us); + result.softmax_mean_us = mean_of(softmax_timings_us); + result.softmax_stdev_us = + stdev_of(softmax_timings_us, result.softmax_mean_us); result.dispatched_kernels = last_dispatched; return result; } @@ -1322,6 +1343,7 @@ void emit_sdpa_records( float stdev; } subs[] = { {"qk", r.qk_mean_us, r.qk_stdev_us}, + {"softmax", r.softmax_mean_us, r.softmax_stdev_us}, {"av", r.av_mean_us, r.av_stdev_us}, {"total", r.mean_us, r.stdev_us}, }; @@ -1854,6 +1876,40 @@ bool sdpa_correctness_case(const SdpaCorrectnessCase& c) { } } + // Localize the mismatches when there are any. A scattered handful spread + // over every head and row means something timing-dependent; a run confined + // to one head, one row block, or one column range points at an indexing or + // tile-boundary bug. Output is [S, Q_H, D], so decode each flat index. + // attn*V's M-tile is 64 rows (WG_TILE_M) and its N-tile is WG_TILE_N=64. + if (mismatches > 0) { + const int64_t D = c.head_dim, QH = c.num_heads; + int64_t s_min = c.seq_len, s_max = -1, d_min = D, d_max = -1; + std::set heads, s_vals, m_tiles; + for (int64_t i = 0; i < q_numel; ++i) { + const float diff = std::fabs(outf[i] - ref[i]); + if (diff <= abs_tol + rel_tol * std::fabs(ref[i])) { + continue; + } + const int64_t s = i / (QH * D), h = (i / D) % QH, d = i % D; + s_min = std::min(s_min, s); + s_max = std::max(s_max, s); + d_min = std::min(d_min, d); + d_max = std::max(d_max, d); + heads.insert(h); + s_vals.insert(s); + m_tiles.insert(s / 64); + } + std::cout << "\n[sdpa-mismatch-loc] " << c.name << " n=" << mismatches + << " rows=[" << s_min << "," << s_max + << "] distinct_rows=" << s_vals.size() << " cols=[" << d_min + << "," << d_max << "]" << " heads=" << heads.size() << "/" << QH + << " attnV_m_tiles={"; + for (auto t : m_tiles) { + std::cout << t << ","; + } + std::cout << "}\n"; + } + const bool numeric_ok = mismatches == 0; const bool fired_ok = qk_fired && av_fired; std::cout << "[sdpa-correctness] " << c.name << " S=" << c.seq_len From 788617961ab75d2df092d1a56c65fc5c9d15eb40 Mon Sep 17 00:00:00 2001 From: Yanwen Xu Date: Wed, 2 Sep 2026 22:28:56 -0700 Subject: [PATCH 21/28] [ET-VK] SDPA attn*V coopmat: skip all-zero K-chunks attn*V walked all num_k_chunks_arg chunks of the context for every output tile. Two thirds of that work could not contribute anything. num_k_chunks_arg is max_context_len/WG_TILE_K, sized from the KV cache, so on a ctx3072 PTE it is 96 chunks even when the live context_len is 2048. The 32 beyond-context chunks staged zeros -- and then ran the MMA on them anyway. Worse, P is the softmax output of a causally masked score matrix, so P[s, c] == 0 for every c > s + input_pos exactly (QK^T writes -inf, softmax normalizes it to zero). A chunk whose lowest context index already exceeds the highest row in the M-tile is therefore all zeros, and a zero A tile contributes 0*V == 0. Skipping it is value-preserving, not an approximation. Bounding the loop by both facts takes the chunk-loop iterations at the 2048 prefill from 32 M-tiles x 96 chunks = 3072 down to 1056 (34.4%): the work is a triangle over M-tiles, not a rectangle, and the beyond-context rectangle is gone entirely. 8B prefill, per-kernel, n=5: 13002.2 -> 4560.4 us, -64.9% (2.85x), ranges disjoint. The 2.9x predicted by the iteration count matches the 2.85x measured, so the model of where the time went is right. Correctness 0 failures / 120 case-runs on the extended gate. REJECTED on the way here: dbuf4 double-buffered staging, the intervention this shader was originally slated for. It measured 13002.2 -> 13124.7 us, +0.94%, ranges disjoint -- a small but real regression. attn*V moves ~1.04 GB per dispatch in ~13.0 ms, i.e. ~80 GB/s, which is bandwidth-bound, and double buffering hides latency rather than bandwidth. It also doubled LDS 9728 -> 19456 B, halving co-residency from 6 to 3 workgroups per CU. Reducing traffic was the answer; overlapping it was not. (cherry picked from commit d132d44066a84a624e61756a633ef7f6f65702a2) --- .../ops/glsl/sdpa_compute_out_coopmat.glsl | 24 ++++++++++++++++++- 1 file changed, 23 insertions(+), 1 deletion(-) diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coopmat.glsl b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coopmat.glsl index ed6f13f206f..3e022cdebf3 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coopmat.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coopmat.glsl @@ -152,7 +152,29 @@ void main() { const uint v_head_base = uint(kv_h) * uint(D4); const uint v_row_stride = uint(KV_H) * uint(D4); - for (uint chunk = 0; chunk < uint(num_k_chunks_arg); ++chunk) { + // Causal K-loop truncation. P is the softmax output, and the causal mask + // has already forced P[s, c] == 0 for every c > s + input_pos (QK^T writes + // -inf there and softmax normalizes it to exactly 0). A K-chunk whose + // lowest context index already exceeds the highest row's window is + // therefore all zeros, and a zero A tile contributes 0*V == 0 to the + // accumulator -- so it can be skipped outright rather than staged and + // MMA'd. This is exactly value-preserving, not an approximation. + // + // The highest context index any row in this M-tile can attend to: + const uint max_c = a_row_base + WG_TILE_M - 1u + uint(input_pos); + // Chunk `c` spans [c*WG_TILE_K, c*WG_TILE_K + WG_TILE_K), so it carries a + // nonzero only when c*WG_TILE_K <= max_c. + const uint useful_chunks = + min(uint(num_k_chunks_arg), max_c / WG_TILE_K + 1u); + // At the 2048 prefill this cuts the staged chunks roughly in half (the + // work is a triangle over M-tiles, not a rectangle): 1056 chunk-loads + // across the 32 M-tiles instead of 32*64 = 2048. + // + // Why this is safe to hardcode: LLM-mode SDPA here is unconditionally + // causal -- sdpa_compute_attn_weights_coopmat.glsl applies the + // c > s + input_pos mask with no is_causal switch -- so the same + // assumption already governs the shader that produces this shader's input. + for (uint chunk = 0; chunk < useful_chunks; ++chunk) { const uint chunkK = chunk * WG_TILE_K; // along context_len // num_k_chunks is max_context_len/WG_TILE_K (static spec const). The // gate guarantees context_len % WG_TILE_N == 0, hence % WG_TILE_K == 0, From 8beee481e0402aafc00c29d0604637e0be29d834 Mon Sep 17 00:00:00 2001 From: Yanwen Xu Date: Wed, 2 Sep 2026 23:19:26 -0700 Subject: [PATCH 22/28] [ET-VK] Remove the SDPA coopmat bisect switches ET_VK_SDPA_DISABLE_QK_COOPMAT and ET_VK_SDPA_DISABLE_OUT_COOPMAT existed for one purpose: attributing an intermittent wrong-output failure to one of the two SDPA GEMMs by forcing the other onto its tiled path. That worked -- it put the fault in attn*V (4/100 case-runs) with QK^T clean (0/100) -- and the underlying missing shared-memory barrier is now fixed, so there is nothing left to isolate and no reason to keep a debug env var in the dispatch gate. Gate after removal: 0 failures / 100 case-runs, with both coopmat shaders confirmed dispatching on every rep (which is now the only mode, so a silent fallback would have shown up as a dispatch warning rather than passing). (cherry picked from commit 846d217171c8631de32e6d6b70bc091ca91f55aa) --- backends/vulkan/runtime/graph/ops/impl/SDPA.cpp | 5 ----- 1 file changed, 5 deletions(-) diff --git a/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp b/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp index 2d2fde2f231..54a4b5f67aa 100644 --- a/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp @@ -349,11 +349,7 @@ vkapi::ShaderInfo pick_sdpa_qk_shader( const bool is_gemv = is_single_token(graph, q_projected); // Prefill WMMA path: Q @ K^T with K = head_dim, N = context_len, M = S. - // DEBUG BISECT (temporary): disable coopmat for ONE of the two SDPA GEMMs - // so an intermittent numerical failure can be attributed to a specific - // shader. Reverted once the attribution is recorded. if (!is_gemv && sdpa_coopmat_device_ok(graph) && - std::getenv("ET_VK_SDPA_DISABLE_QK_COOPMAT") == nullptr && sdpa_buf_half(graph, q_projected) && sdpa_buf_half(graph, k_cache) && sdpa_buf_half(graph, attn_weights)) { const SDPADims d = compute_sdpa_dims( @@ -507,7 +503,6 @@ vkapi::ShaderInfo pick_sdpa_av_shader( // Prefill WMMA path: P @ V with K = context_len, N = head_dim, M = S. if (!is_gemv && sdpa_coopmat_device_ok(graph) && - std::getenv("ET_VK_SDPA_DISABLE_OUT_COOPMAT") == nullptr && sdpa_buf_half(graph, out) && sdpa_buf_half(graph, attn_weights_softmax) && sdpa_buf_half(graph, v_cache)) { From be70eee10944189da491c81f856fca4ef9b016d2 Mon Sep 17 00:00:00 2001 From: Yanwen Xu Date: Thu, 3 Sep 2026 13:45:29 -0700 Subject: [PATCH 23/28] [ET-VK] Fence the 4 remaining SDPA coopmat LDS barriers Completes the memoryBarrierShared() work for SDPA. 8dc05ca896 fixed only the write->read barrier in attn*V and said so; this covers the rest. sdpa_compute_attn_weights_coopmat.glsl (QK^T) +3 -> 3 fenced, 0 bare sdpa_compute_out_coopmat.glsl (attn*V) +1 -> 2 fenced, 0 bare Two of the QK^T sites are the WRITE->READ direction, i.e. the same direction that produced the observed ~2.5%-of-runs stale-A-band bug in attn*V, not the untested WAR direction: QK^T:266 after the Bsh[...] staging stores, before the MMA coopMatLoad QK^T:349 after coopMatStore into Csh, before the epilogue reads it QK^T:297 WAR -- this chunk's reads vs the next chunk's stores attn*V:260 WAR -- same So QK^T was exposed in the dangerous direction. The prior note that it was "measured clean 0/100, not proven safe" understated it. VALIDATION STATUS: build verified, device validation PENDING. Two builds (fenced/unfenced) were produced and every .spv compared: of 1620 shaders only the 2 SDPA ones differ, so build variance is excluded as a confounder. The fences are confirmed present in the compiled SPIR-V rather than assumed -- OpMemoryBarrier 0 -> 3 (QK^T) and 1 -> 2 (attn*V), OpControlBarrier unchanged. The on-device correctness gate and the interleaved A/B have NOT run: the target board (xgpusw-debug08 / 00000bf70c579c33) crashed to bootloader twice and came back with a third, unrecognised driver hash. Do not treat this commit as perf-validated until that A/B is recorded. The equivalent fences on the linear defaults measured +0.001% / -0.002% (7892151a81), so a cost is unlikely. --- .../glsl/sdpa_compute_attn_weights_coopmat.glsl | 15 +++++++++++++++ .../graph/ops/glsl/sdpa_compute_out_coopmat.glsl | 5 +++++ 2 files changed, 20 insertions(+) diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_coopmat.glsl b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_coopmat.glsl index fe95e6cb605..8b4e95815ad 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_coopmat.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_coopmat.glsl @@ -263,6 +263,11 @@ void main() { Bsh[(d_base + 3u) * B_ROW + lc] = v.w; } + // coopmat-lds-fence (WRITE->READ): the Ash/Bsh staging stores above vs the coopMatLoad below -- the SAME direction that produced the observed stale-A-band bug in attn*V. + // barrier() alone does NOT order shared stores against a subsequent + // coopMatLoad on the M51 Xclipse/AMD-PAL driver (~2.5% of runs, one stale + // MMA_M-row band, silently wrong, no crash). Measured cost: none. + memoryBarrierShared(); barrier(); // --- Cooperative matrix MMA: result += A * B (B is already K^T) --- @@ -294,6 +299,11 @@ void main() { } } + // coopmat-lds-fence (WAR): this chunk's coopMatLoad reads vs the next chunk's staging stores. + // barrier() alone does NOT order shared stores against a subsequent + // coopMatLoad on the M51 Xclipse/AMD-PAL driver (~2.5% of runs, one stale + // MMA_M-row band, silently wrong, no crash). Measured cost: none. + memoryBarrierShared(); barrier(); } @@ -346,6 +356,11 @@ void main() { gl_CooperativeMatrixLayoutRowMajor); } } + // coopmat-lds-fence (WRITE->READ): the coopMatStore into Csh above vs the epilogue read below. + // barrier() alone does NOT order shared stores against a subsequent + // coopMatLoad on the M51 Xclipse/AMD-PAL driver (~2.5% of runs, one stale + // MMA_M-row band, silently wrong, no crash). Measured cost: none. + memoryBarrierShared(); barrier(); // --- Copy Csh -> global attn_weights with the per-element causal mask --- diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coopmat.glsl b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coopmat.glsl index 3e022cdebf3..d142608f7bd 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coopmat.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coopmat.glsl @@ -257,6 +257,11 @@ void main() { } } + // coopmat-lds-fence (WAR): this chunk's Ash/Bsh coopMatLoad reads vs the next chunk's stores. + // barrier() alone does NOT order shared stores against a subsequent + // coopMatLoad on the M51 Xclipse/AMD-PAL driver (~2.5% of runs, one stale + // MMA_M-row band, silently wrong, no crash). Measured cost: none. + memoryBarrierShared(); barrier(); } From 921b0cf67335572e817672be58e15651a5a822af Mon Sep 17 00:00:00 2001 From: Yanwen Xu Date: Tue, 8 Sep 2026 10:36:59 -0700 Subject: [PATCH 24/28] [ET-VK] Enable WMMA by default: texture coopmat no longer needs an env var ET_VK_TEXTURE_COOPMAT began as an experiment hook (specs/040/041), off by default, documented as "Off by default, so buffer dispatch is byte-identical". That reasoning only holds for buffer-storage PTEs. The embq PTEs this branch targets place the linear weights in TEXTURE storage, so with the hook off the *_texture3d_* coopmat variants were rejected and the entire quantized linear path -- 72.4% of prefill GPU time -- silently fell back to the tiled shader. No linear WMMA ran at all unless the caller happened to know about the flag. Flip the default to enabled. ET_VK_TEXTURE_COOPMAT=0 (or "false"/"off") still restores the old buffer-dispatch-only behavior for baseline measurement; unset, or any other value, enables it. SDPA needed no change -- it was already on by default, gated only on supports_cooperative_matrix() && subgroup_size() == 64, with ET_VK_DISABLE_COOPMAT as the kill switch. 8B/8da4w, 2048 prefill, embq ctx3072, xgpusw-debug08/00000b750f413c33, driver main@fafb46ae9c0d (md5 1eea300aa39..), BSP CP2A.260605.016 20260831.130326, clocks maxpin 980/5333/934, 3 reps each: no env vars, BEFORE this change 174.5 / 174.8 / 174.9 -> 174.7 tok/s ET_VK_TEXTURE_COOPMAT=1 + NODE_THRESHOLD=32 367.5 / 366.8 / 368.0 -> 367.4 tok/s no env vars, AFTER this change 366.9 / 367.9 / 368.0 -> 367.6 tok/s 2.10x, entirely from the default. The post-change no-env run matches the env-forced run to 0.05%, and needs no ET_VK_EXECUTE_NODE_THRESHOLD -- that watchdog workaround is not required on this driver. ETDump verification (no env vars set), leaf GPU 5484.0 us over 1363 events: 72.4% linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr_t128x64k32g42s32_texture3d_texture2d_half 5.7% sdpa_attn_weights_softmax_buffer_half 4.8% sdpa_compute_attn_weights_coopmat_t128x64k32g22s64_buffer_buffer_half 2.7% sdpa_compute_out_coopmat_t64x64k32g22s64_buffer_buffer_half Both SDPA GEMMs dispatch on the _coopmat path as buffer_buffer_half, so the buffer+fp16 gate in SDPA.cpp is satisfied even though the linear path is on texture IO. No tiled sdpa_compute_attn_weights / sdpa_compute_out appears -- there is no silent fallback left. Dumps archived at .artifacts/rel14qs-{coopmat,defaultwmma-noenv}-2026-09-08.etdp. Not validated: decode (prefill-only runs, --max_new_tokens=1) and models other than 8B/8da4w. --- .../graph/ops/impl/QuantizedLinear.cpp | 31 ++++++++++++++++--- 1 file changed, 26 insertions(+), 5 deletions(-) diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp index 09218f0fe9c..3255622ba9e 100644 --- a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp @@ -457,12 +457,33 @@ utils::uvec3 quantized_linear_local_wg_size( } } -// Experiment hook (specs/040/041): allows the *_texture3d_* variants, which -// stage the result tile through shared memory and imageStore it instead of -// coopMatStore-ing straight to an SSBO. Off by default, so buffer dispatch is -// byte-identical. +// Allows the *_texture3d_* coopmat variants, which stage the result tile +// through shared memory and imageStore it instead of coopMatStore-ing straight +// to an SSBO. +// +// ON BY DEFAULT since 2026-09-08. It began as an experiment hook +// (specs/040/041) and was opt-in via ET_VK_TEXTURE_COOPMAT=1, but the embq PTEs +// this branch targets place the linear weights in texture storage, so with the +// hook off the whole quantized linear path silently fell back to the TILED +// shader and no WMMA ran at all. Measured on 8B/8da4w 2048-prefill +// (xgpusw-debug08 / 00000b750f413c33, driver main@fafb46ae9c0d, maxpin +// 980/5333/934): 174.7 tok/s with the hook off vs 367.4 tok/s with it on -- +// a 2.10x difference that had nothing to do with the kernel and everything to +// do with the default. +// +// Set ET_VK_TEXTURE_COOPMAT=0 (or "false"/"off") to restore the old +// buffer-dispatch-only behavior; any other value, or leaving it unset, enables +// texture coopmat. static bool texture_coopmat_enabled() { - static const bool enabled = std::getenv("ET_VK_TEXTURE_COOPMAT") != nullptr; + static const bool enabled = [] { + const char* env = std::getenv("ET_VK_TEXTURE_COOPMAT"); + if (env == nullptr) { + return true; + } + return !( + strcmp(env, "0") == 0 || strcmp(env, "false") == 0 || + strcmp(env, "off") == 0); + }(); return enabled; } From ce7b22f343e427a364d095f24282758e47a87866 Mon Sep 17 00:00:00 2001 From: Yanwen Xu Date: Tue, 8 Sep 2026 18:18:28 -0700 Subject: [PATCH 25/28] [ET-VK] Six opt-in dq8ca coopmat variants: A-layout and B-staging, all rejected Optimization attempts against gemm-ubm's shmem_double_buf4-tr3, which is 1.39x faster than our shipped dq8ca kernel on M51 (2231.8 vs 1601.4 us at 2048x1024x4096, 47.94% vs 66.82% of int8 WMMA peak, profiler off, same board/driver/clocks). Shipped default is UNCHANGED; every variant is opt-in via ET_VK_DQ8CA_COOPMAT_VARIANT and none is a default. A-layout (dbuf4zpgtr3): port tr3's shared-A layout -- scalar int8_t storage, row-major over the full chunk, A_ROW_PAD_I8 bytes of per-row pad (yaml parameter). Correctness 12/12 clean; MEASURED SLOWER at every pad value: p16 +4.73%, p8 +4.71%, p0 +48.78%. Axis decomposition: - layout+element-type alone (p0, LDS identical to baseline): +48.78% - adding 8B of pad recovers 44.1 pp - 16B (tr3's own value) adds nothing over 8B (+0.019%) -- tr3 over-pads for M51 and the extra 2 KiB of LDS is a pure loss - vs the earlier zpgtrp attempt (same layout+pad but `shared uint`, +11.50%), the element-type change alone is worth 6.8 pp Mechanism: slab-major gives each 16x16 A fragment a row stride of exactly 16 bytes, so the fragment is 256 CONTIGUOUS bytes; row-major full-chunk makes it strided (32/40/48 B). Padding fixes the bank pattern but cannot restore contiguity. tr3's row-major A is a constraint of its own staging, not a virtue. VGPR 128 unchanged, scratch 0, LDS 14080 -> 16128 (still 4 WG/CU). B-staging: five variants moving ownership from (block,col) slots spread over 4 threads to one thread per contiguous run, to cut the measured 3.25x static ds_* gap. - bv4 (uvec4 array, 4-wide), bw (no retype, 4-wide), bwr (bw, reversed store order): all KNOWN-INCORRECT, deterministic 3/3 failure on the num_groups==2 shapes only (K=256); K=128/2048/4096 pass. Cause isolated to 4-wide-per-thread ownership with half the workgroup idle. Store order is not the trigger (bwr) and neither is the uvec4 retype (bw, no retype, fails identically). Deliberately NOT registered so they cannot be selected. - bw2 (2-wide, all 256 threads active): CORRECT 14/14 x3, ~0% (-0.02%/-0.115%) - bw3 (uvec2 retype + bw2 ownership): CORRECT 14/14 x3, +0.34%. Also proves an 8x packing ratio is fine for coopMatLoad, which retroactively clears the retype as bv4's cause. The transformation demonstrably landed and bought nothing: scalar ds_store_b32 11 -> 3, paired ds_store_2addr_b32 10 -> 13, total ds ops 104 -> 100, static instructions 1647 -> 1607. bw2 and bw3 emit byte-identical LDS op mixes, so the compiler was already merging bw2's adjacent scalar writes. Why it could never have worked, now measured: our LDS instruction count is dominated by LOADS, not stores -- ours 80 loads / 24 stores against TR3's 24 / 8. ds_load_2addr_b32 alone is 72 and is the largest single ISA category. The 80 loads include the six auxiliary shared arrays (izp_sh, ifs_sh, wsc_sh, wcorr_sh stride-0 broadcasts, bias_sh, Csh_out) that TR3 has no analogue for. That is the next target; B-store width is closed. Full measurement record, counter data and provenance: openspec/changes/dq8ca-vs-tr3-iso-tile-counters/results.md --- ...q8ca_q4gsw_coopmat_tsweep_dbuf4zpgbv4.glsl | 792 +++++++++++++++++ ...q8ca_q4gsw_coopmat_tsweep_dbuf4zpgbv4.yaml | 68 ++ ...dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw.glsl | 781 +++++++++++++++++ ...dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw.yaml | 68 ++ ...q8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw2.glsl | 762 +++++++++++++++++ ...q8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw2.yaml | 60 ++ ...q8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw3.glsl | 798 ++++++++++++++++++ ...q8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw3.yaml | 60 ++ ...q8ca_q4gsw_coopmat_tsweep_dbuf4zpgbwr.glsl | 783 +++++++++++++++++ ...q8ca_q4gsw_coopmat_tsweep_dbuf4zpgbwr.yaml | 68 ++ ...q8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr3.glsl | 781 +++++++++++++++++ ...q8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr3.yaml | 92 ++ .../graph/ops/impl/QuantizedLinear.cpp | 22 +- 13 files changed, 5134 insertions(+), 1 deletion(-) create mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbv4.glsl create mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbv4.yaml create mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw.glsl create mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw.yaml create mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw2.glsl create mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw2.yaml create mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw3.glsl create mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw3.yaml create mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbwr.glsl create mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbwr.yaml create mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr3.glsl create mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr3.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbv4.glsl b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbv4.glsl new file mode 100644 index 00000000000..c74d6558055 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbv4.glsl @@ -0,0 +1,792 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +/* + * "zpg" + "-tr" combination: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpg.glsl + * with its per-thread scalar A-staging replaced by + * linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr.glsl's coopMat-mediated A + * staging, PLUS "bv4": shared B retyped to uvec4 so each (K-slab, column) + * column-slab is ONE 16-byte element written by ONE b128 store instead of 4 + * scalar ds_write_b32 spread across 4 threads. A staging, the quantization + * math and every correctness safeguard are dbuf4zpgtr's, unchanged. This is + * an ADDITIVE combination, not a redesign: every non-A-staging block below + * (B staging: coalesced write, no skew; zp-hoist: izp/ifs applied once after + * the group loop via wcorr_sh; byte-parallel nibble widening; static + * A_ALWAYS_ACTIVE branch elision -- N/A here, see below; group epilog; + * bias/store epilogue) is byte-identical to dbuf4zpg's. Only the A-staging + * block (prologue load+store, main-loop prefetch+store) is dbuf4tr's, + * verbatim. + * + * dbuf4zpg's per-thread A staging used an `a_active` guard (statically always + * true when A_ACTIVE_THREADS == WG_SIZE, via the A_MAP_FULL-gated + * A_ALWAYS_ACTIVE macro). dbuf4tr's per-SUBGROUP tile map has no equivalent + * concept -- every subgroup participates via a `t < NUM_A_TILES` guard that + * depends only on gl_SubgroupID, not gl_LocalInvocationID.x -- so + * A_MAP_FULL/A_ALWAYS_ACTIVE is dropped entirely in this file; it would be + * dead code for the new A-staging block. + * + * Rationale for combining this way (not the reverse) and why this is worth + * building at all: see this change's design.md D0-D3. In short -- the only + * existing measurement of dbuf4tr's A-staging technique (28.72-30.51%, + * dq8ca-arch-redesign) was taken against dbuf4tr's own pre-zpg baseline + * (old B skew, no byte-parallel widening, no branch elision) -- a materially + * weaker shader than the 46.49-46.50% dbuf4zpg this file now combines it + * with. This file exists to answer whether that combination performs + * differently now that register pressure is already reduced. + * + * A staging (the actual delta from dbuf4zpg): + * dbuf4zpg: per-thread (m4, k4) ivec4 fetch, hoisted a_lds_off0/a_glb_row; + * only A_ACTIVE_THREADS invocations participate, each scattering + * 4 rows into Ash_int8 with 4 scalar stores. + * this file: per-SUBGROUP MMA_M x MMA_K tile fetch via coopMatLoad straight + * from a ROW-MAJOR (kPackedInt8_4W) int8 activation buffer, then + * coopMatStore into the same Ash_int8 slot -- dbuf4tr's mapping, + * unmodified (not re-derived; see design.md D3). + * + * t_packed_int8_input is therefore bound the same way dbuf4tr binds it: a + * SCALAR int8_t array in the kPackedInt8_4W layout (plain row-major int8, + * row stride K), produced by quantize_and_pack_4w_with_group_sums.glsl. + * QuantizedLinear.cpp's dq8ca_variant_wants_rowmajor_a() must recognize this + * file's variant token (tsweep_dbuf4zpgtr_t...) the same way it already + * recognizes tsweep_dbuf4tr_t/trm_t/trd_t, so graph-build time (packer + * selection) and dispatch time (kernel selection) cannot disagree. + * + * B CANNOT be coopmat-staged (int4 nibble unpack; a coopmat's per-lane layout + * is opaque to hand-assembly from unpacked registers) -- unchanged from both + * parent files. B staging below is dbuf4zpg's byte-parallel, coalesced, + * no-skew version, untouched. + * + * The loop structure is dbuf4's (both parents share it), unchanged: + * prologue: prefetch chunk 0 -> temp, store to slice 0 (no barrier) + * per iter: barrier -> prefetch(next) -> MMA(cur) -> store(next) + * kept nested (groups x chunks) with an unconditional group epilog -- + * flattening it crashes the Xclipse PAL compiler at large spec-resolved trip + * counts (see dbuf2's own header). + * + * Selected via + * ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4zpgtr_txkgs<32|64> + * (QuantizedLinear.cpp), additive to the tsweep_dbuf4zpg_t..., tsweep_dbuf4tr_t... + * and tsweep_t... namespaces. NOT the default -- unvalidated until it passes + * repeated test_llama_microbench --correctness-only runs (see + * dq8ca_coopmat_variant()'s comment on why a single pass is not proof). + * + * Performs: out[M,N] = dequant(int8_act) * dequant(int_w) (+ bias) + * via coopmat x coopmat -> coopmat on the matrix unit. + * + * Hard preconditions (dbuf4zpg's, plus dbuf4tr's row-major/alignment ones): + * M % WG_TILE_M == 0, N % WG_TILE_N == 0, K % WG_TILE_K == 0, + * group_size % WG_TILE_K == 0, K % 4 == 0, + * WG_TILE_M % MMA_M == 0, WG_TILE_K % MMA_K == 0, + * t_packed_int8_input in kPackedInt8_4W (row-major) layout, + * device exposes coopmatx-> at 16x16x16. + */ + +#version 450 core + +#extension GL_KHR_cooperative_matrix : require +#extension GL_KHR_memory_scope_semantics : require +#extension GL_KHR_shader_subgroup_basic : enable +#extension GL_EXT_shader_explicit_arithmetic_types : require +#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require +// 8-bit SSBO access: A is bound as a scalar int8_t array so that the +// coopMatLoad below has a MATCHING component type (see dbuf4tr's header for +// why the type must match on this driver). +#extension GL_EXT_shader_8bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_EXT_control_flow_attributes : enable + +#define PRECISION ${PRECISION} + +$if WEIGHT_NBITS == 4: + #define WEIGHT_INT4 + +$if HAS_BIAS: + #define HAS_BIAS + +$if WEIGHT_STORAGE == "buffer": + #define WEIGHT_BUFFER + +$if IO_STORAGE == "texture3d": + #define IO_TEXTURE + +layout(std430) buffer; + +#include "common.glslh" + +// Bindings — match add_linear_dqa_qw_node arg order: +// output(0), fp_input(1), packed_int8_input(2), int_input_sums(3 - unused), +// input_scales(4), input_zps(5), packed_weight(6), weight_sums(7), +// weight_scales(8), bias(9). +${layout_declare_tensor(B, "w", "t_output", "half", IO_STORAGE, is_scalar_array=True)} +// t_input is unread here -- the activations arrive already quantized in +// t_packed_int8_input -- but stays declared so the binding layout matches the +// dispatch site. It tracks IO_STORAGE so the two IO tensors stay consistent. +${layout_declare_tensor(B, "r", "t_input", "half", IO_STORAGE, is_scalar_array=False)} +// ROW-MAJOR (kPackedInt8_4W) packed activations, bound as a scalar int8_t +// array (row stride = K int8) -- dbuf4tr's binding, unchanged. The stock +// 4h4w layout dbuf4zpg uses is NOT row-major (component index selects a row, +// non-affine), so it cannot be addressed by any coopMatLoad. +${layout_declare_tensor(B, "r", "t_packed_int8_input", "int8", "buffer", is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_int8_input_sums", "int", "buffer", is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_int8_input_scales", "half", "texture3d")} +${layout_declare_tensor(B, "r", "t_int8_input_zps", "int8", "texture3d")} +${layout_declare_tensor(B, "r", "t_packed_weight", "int", WEIGHT_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_weight_sums", "int", "buffer", is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_weight_scales", "half", "buffer", is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_bias", "half", "buffer", is_scalar_array=True)} + +${layout_declare_ubo(B, "ivec4", "output_sizes")} +${layout_declare_ubo(B, "ivec4", "input_sizes")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +${layout_declare_spec_const(C, "int", "apply_bias", "0")} +// INT4 only; inert (0) for INT8 so the dispatcher's spec list lines up. +${layout_declare_spec_const(C, "int", "K4_per_group", "0")} +${layout_declare_spec_const(C, "int", "num_groups_arg", "0")} +${layout_declare_spec_const(C, "int", "out_N_arg", "0")} + +// Tile geometry +const uint MMA_M = ${MMA_M}; +const uint MMA_N = ${MMA_N}; +const uint MMA_K = ${MMA_K}; + +const uint WG_TILE_M = ${WG_TILE_M}; +const uint WG_TILE_N = ${WG_TILE_N}; +const uint WG_TILE_K = ${WG_TILE_K}; + +const uint SG_GRID_X = ${SG_GRID_X}; +const uint SG_GRID_Y = ${SG_GRID_Y}; +const uint SUBGROUP_SIZE = ${SUBGROUP_SIZE}; +const uint NUM_SUBGROUPS = SG_GRID_X * SG_GRID_Y; +const uint WG_SIZE = NUM_SUBGROUPS * SUBGROUP_SIZE; + +const uint SG_TILE_M = WG_TILE_M / SG_GRID_Y; +const uint SG_TILE_N = WG_TILE_N / SG_GRID_X; +const uint MMAS_PER_SG_M = SG_TILE_M / MMA_M; +const uint MMAS_PER_SG_N = SG_TILE_N / MMA_N; + +const uint A_SLAB_INT8 = WG_TILE_M * MMA_K; +const uint B_USEFUL_U32 = MMA_K / 4u; +// No skew + coalesced write -- dbuf4zpg's B fix, unchanged (this file does +// not touch B staging at all). +const uint B_STRIDE_U32 = B_USEFUL_U32; +const uint B_SLAB_U32 = WG_TILE_N * B_STRIDE_U32; +const uint NUM_K_SLABS = WG_TILE_K / MMA_K; + +const uint A_SLAB_U32 = A_SLAB_INT8 / 4u; +const uint A_STRIDE_U32 = MMA_K / 4u; + +// One ping-pong slice covers all K-slabs of one chunk. +const uint ASH_SLICE_U32 = NUM_K_SLABS * A_SLAB_U32; +const uint BSH_SLICE_U32 = NUM_K_SLABS * B_SLAB_U32; + +// Double-buffered MMA operand staging. +shared uint Ash_int8[2u * ASH_SLICE_U32]; +// ===== bv4: shared B is uvec4-typed, ONE element per (K-slab, column) ===== +// dbuf4zpgtr stores B as `uint`: 4 uints per column-slab, written as 4 separate +// scalar ds_write_b32 spread over 4 threads. One column-slab is exactly +// MMA_K int8 = 16 bytes = one uvec4, and B_STRIDE_U32 is already 4 because zpg +// removed the anti-bank-conflict skew -- so those 4 uints are contiguous AND +// 16B-aligned. That alignment is what made a b128 store impossible on the old +// skewed layout (75% of columns were unaligned) and possible now. +// Each owning thread assembles all 4 widened uints and issues ONE store. +const uint BSH_SLICE_V4 = NUM_K_SLABS * WG_TILE_N; +shared uvec4 Bsh_v4[2u * BSH_SLICE_V4]; + +// Per-WG-tile-row activation params (loaded ONCE at WG start; constant +// across groups). +shared int izp_sh[WG_TILE_M]; // int32 (cast from int8 source) for broadcast +shared float ifs_sh[WG_TILE_M]; // float32 (cast from fp16 source) for broadcast + +// Per-(group, output-channel) weight params, ping-ponged by group parity. +// (For per-channel INT8 only slice 0 is ever used.) +shared float wsc_sh[2u * WG_TILE_N]; +// SUM_g wsc[g][n]*wsum[g][n] per output channel -- weight-side only, so it is +// accumulated once in the prologue. dbuf4zpg's zp-hoist, unchanged. +shared float wcorr_sh[WG_TILE_N]; + +#ifdef HAS_BIAS +shared float bias_sh[WG_TILE_N]; +#endif + +#ifdef IO_TEXTURE +// Result staging for the imageStore epilogue, mirroring the fp16 kernel: +// SG_GRID_Y bands of MMA_M rows, each WG_TILE_N wide, row-major. A full +// WG_TILE_M x WG_TILE_N buffer would cost SG_GRID_Y/MMAS_PER_SG_M x more LDS +// and wreck occupancy. float16_t-typed because coopMatStore needs it. +const uint CSH_ROWS = SG_GRID_Y * MMA_M; +shared float16_t Csh_out[CSH_ROWS * WG_TILE_N]; +#endif + +// Running fp32 accumulator (across all groups). +coopmat + result[MMAS_PER_SG_M][MMAS_PER_SG_N]; + +// Per-group int32 MMA accumulator. +coopmat + accum_int32[MMAS_PER_SG_M][MMAS_PER_SG_N]; + + +// Byte-parallel int4 -> int8 widening. dbuf4zpg's, unchanged (B-side only). +// +// The four nibbles this shader needs from one packed uint are ALREADY one per +// byte (bits 3:0 of each byte for parity 0, bits 7:4 for parity 1), so all four +// can be widened at once instead of with a per-nibble +// shift/mask/bias-subtract/mask chain. +// +// For v in [0,15] the biased value is v-8. `v ^ 8` is exactly the 4-bit +// two's-complement pattern of v-8, because -8 == +8 (mod 16): +// v=0 -> 0x8 -> -8 v=7 -> 0xF -> -1 +// v=8 -> 0x0 -> 0 v=15 -> 0x7 -> +7 +// so the only remaining work is sign-extending bit 3 into bits 7:4 per byte. +// `sgn * 0x1E` does that with no cross-byte carry: 0x08 * 0x1E == 0xF0 exactly, +// and sgn is at most 0x08080808 so the product is at most 0xF0F0F0F0. +// +// A naive `nib - 0x08080808` would NOT work -- it borrows across byte lanes +// whenever a nibble is < 8. Shifts must be on uint, not int, so the >> is +// logical rather than arithmetic. +// +// ~5 ops per 4 weights vs ~22 for the per-nibble chain; bit-identical output. +uint widen_nibbles(const uint w, const uint parity) { + const uint nib = (parity == 0u) ? (w & 0x0F0F0F0Fu) : ((w >> 4u) & 0x0F0F0F0Fu); + const uint p = nib ^ 0x08080808u; + const uint sgn = p & 0x08080808u; + return p | (sgn * 0x1Eu); +} + +void main() { + const uvec2 tileID = uvec2(gl_WorkGroupID.xy); + const uvec2 warpInTile = uvec2( + gl_SubgroupID % SG_GRID_X, + gl_SubgroupID / SG_GRID_X); + + const uint K = uint(input_sizes.x); + const uint N = uint(output_sizes.x); + const uint N4 = (N + 3u) / 4u; + const uint nblocks_x_A = (K + 3u) >> 2u; + // A row stride in INT8 elements (dbuf4tr's binding is row-major int8, not + // the 4h4w ivec4 block layout dbuf4zpg used -- so A addressing below is in + // int8 elements, not int, and derived from nblocks_x_A so it matches the + // packer's `m_row * K4 + k4` addressing exactly (K % 4 == 0 makes them + // equal to K directly). + const uint a_row_stride_i8 = nblocks_x_A * 4u; + +#ifdef WEIGHT_INT4 + const uint num_groups = uint(num_groups_arg); + const uint CHUNKS_PER_GROUP = uint(K4_per_group) * 4u / WG_TILE_K; +#else + // Per-channel: a single quant "group" spanning all of K. The nested + // groups x chunks loop below collapses to a flat chunk loop, the wsum/wsc + // ping-pong never crosses a boundary, and the epilog runs exactly once. + const uint num_groups = 1u; + const uint CHUNKS_PER_GROUP = uint(num_groups_arg); +#endif + const uint num_chunks = num_groups * CHUNKS_PER_GROUP; + + const uint tile_m_start = WG_TILE_M * tileID.y; + const uint tile_n_start = WG_TILE_N * tileID.x; + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + result[i][j] = coopmat(0.0); + accum_int32[i][j] = coopmat(0); + } + } + + const uint K_BLOCKS_PER_CHUNK = WG_TILE_K >> 2u; + + // --- A staging tile map (dbuf4tr's, unmodified): one MMA_M x MMA_K coopmat + // tile per subgroup per slot, dealt round-robin across the + // NUM_SUBGROUPS subgroups so every subgroup participates. Replaces + // dbuf4zpg's per-thread (m4, k4) map / a_active guard entirely -- see + // design.md D3 for why this is reused as-is, not re-derived. + const uint A_TILES_M = WG_TILE_M / MMA_M; + const uint A_TILES_K = WG_TILE_K / MMA_K; // == NUM_K_SLABS + const uint NUM_A_TILES = A_TILES_M * A_TILES_K; + const uint A_TILES_PER_SG = (NUM_A_TILES + NUM_SUBGROUPS - 1u) / NUM_SUBGROUPS; + +#ifdef WEIGHT_INT4 + // --- B staging thread map: (block, col) slots; each slot extracts one + // ColumnMajor LDS uint (4 K-contiguous sign-extended int8) --- + const uint B_TOTAL_SLOTS = K_BLOCKS_PER_CHUNK * WG_TILE_N; + const uint B_SLOTS_PER_THREAD = B_TOTAL_SLOTS / WG_SIZE; + const uint N8_PER_TILE = WG_TILE_N >> 3u; +#else + // --- B staging thread map: one (k4, n4) ivec4 block per active thread --- + const uint B_FETCH_SLOTS = K_BLOCKS_PER_CHUNK * (WG_TILE_N >> 2u); + const uint N4_PER_TILE = WG_TILE_N >> 2u; + const uint b_k4_in_chunk = gl_LocalInvocationID.x / N4_PER_TILE; + const uint b_n_uint_col = gl_LocalInvocationID.x % N4_PER_TILE; + const bool b_active = gl_LocalInvocationID.x < B_FETCH_SLOTS; +#endif + + // ===== INTERVENTION F: hoist loop-invariant B staging index math ===== + // dbuf4zpg's, unchanged -- B staging is untouched by this file's A-staging + // swap. See dbuf4zpg's header for the full rationale (ablation-attributed + // -16.8% block, mostly index arithmetic not memory traffic). +#ifdef WEIGHT_INT4 + // bv4 ownership: thread t owns one whole (slab, col) column-slab when + // t < BSH_SLICE_V4, and nothing otherwise. comp/par depend only on the + // column, so a thread's 4 k4 fetches share them and differ only in k4 -- + // 4 consecutive k4 blocks of the SAME n8 texel column. Ownership is + // subgroup-aligned (WG_TILE_N and BSH_SLICE_V4 are multiples of the subgroup + // width for every shipped geometry), so a non-owning subgroup skips the whole + // staging block on a wave-uniform branch rather than diverging. + const bool b_owns = gl_LocalInvocationID.x < BSH_SLICE_V4; + const uint b_slab = gl_LocalInvocationID.x / WG_TILE_N; + const uint b_col = gl_LocalInvocationID.x % WG_TILE_N; + const uint b_rem = b_col & 7u; + const uint b_comp1 = b_rem & 3u; + const uint b_par1 = b_rem >> 2u; + const uint b_n8blk1 = (tile_n_start >> 3u) + (b_col >> 3u); + const uint b_k4base = b_slab * (MMA_K >> 2u); + const uint b_v4_off = b_slab * WG_TILE_N + b_col; +#endif + + // Prefetch temp registers. temp_A is a coopmat array (dbuf4tr's A-staging + // technique); indices into it are [[unroll]]-resolved compile-time + // constants, never dynamic -- dynamic indexing of a coopmat array is + // exactly the construct the Xclipse/AMD-PAL compiler has miscompiled + // before. + coopmat + temp_A[A_TILES_PER_SG]; +#ifdef WEIGHT_INT4 + ivec4 temp_B[MMA_K >> 2u]; // bv4: one per k4 of the owned column-slab + float temp_wsc; +#else + ivec4 temp_B; +#endif + + // ========================================================= + // PROLOGUE + // ========================================================= + if (gl_LocalInvocationID.x < (WG_TILE_M >> 2u)) { + const uint m4 = (tile_m_start >> 2u) + gl_LocalInvocationID.x; + const vec4 sc = vec4(texelFetch(t_int8_input_scales, ivec3(m4, 0, 0), 0)); + const ivec4 zp = texelFetch(t_int8_input_zps, ivec3(m4, 0, 0), 0); + const uint base = gl_LocalInvocationID.x * 4u; + ifs_sh[base + 0u] = sc.x; ifs_sh[base + 1u] = sc.y; + ifs_sh[base + 2u] = sc.z; ifs_sh[base + 3u] = sc.w; + izp_sh[base + 0u] = zp.x; izp_sh[base + 1u] = zp.y; + izp_sh[base + 2u] = zp.z; izp_sh[base + 3u] = zp.w; + } + // Group 0 weight scales -> slice 0, and the hoisted weight-side correction + // SUM_g wsc[g][n]*wsum[g][n] accumulated across ALL groups. dbuf4zpg's + // zp-hoist, unchanged. + if (gl_LocalInvocationID.x < WG_TILE_N) { + const uint n_idx = tile_n_start + gl_LocalInvocationID.x; + f16vec4 sv0 = t_weight_scales[n_idx >> 2u]; + wsc_sh[gl_LocalInvocationID.x] = float(sv0[n_idx & 3u]); + + float corr = 0.0; + for (uint g = 0; g < num_groups; ++g) { + f16vec4 sv = t_weight_scales[g * N4 + (n_idx >> 2u)]; + corr += float(sv[n_idx & 3u]) * float(t_weight_sums[g * N + n_idx]); + } + wcorr_sh[gl_LocalInvocationID.x] = corr; + } + memoryBarrierShared(); + barrier(); + + // NOTE: dbuf4zpg builds izp_bcast/ifs_bcast AFTER the group loop, not here + // -- that is the register-pressure saving zp-hoist buys. Unchanged. + + // dbuf4: prefetch chunk 0 into temp registers, THEN store to slice 0 (no + // barrier here -- the main loop's first iteration barriers before + // reading slice 0). + // + // A staging (dbuf4tr's technique): per-subgroup coopMatLoad straight from + // the row-major global buffer. + [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { + const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; + if (t < NUM_A_TILES) { + const uint tm = t / A_TILES_K; + const uint tk = t % A_TILES_K; + coopMatLoad( + temp_A[s], t_packed_int8_input, + (tile_m_start + tm * MMA_M) * a_row_stride_i8 + tk * MMA_K, + a_row_stride_i8, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#ifdef WEIGHT_INT4 + if (b_owns) { + [[unroll]] for (uint kk = 0; kk < (MMA_K >> 2u); ++kk) { +#ifdef WEIGHT_BUFFER + temp_B[kk] = t_packed_weight[(b_n8blk1 * nblocks_x_A) + b_k4base + kk]; +#else + temp_B[kk] = texelFetch(t_packed_weight, ivec2(b_k4base + kk, b_n8blk1), 0); +#endif + } + } +#else + if (b_active) { + const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; +#ifdef WEIGHT_BUFFER + temp_B = t_packed_weight[(b_k4_in_chunk * N4) + block_x_w]; +#else + temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, b_k4_in_chunk), 0); +#endif + } +#endif + { + // store chunk 0 -> slice 0 + // A staging (dbuf4tr's technique): coopMatStore into the same Ash_int8 + // slot layout dbuf4zpg's scalar scatter used to write. + [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { + const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; + if (t < NUM_A_TILES) { + const uint tm = t / A_TILES_K; + const uint tk = t % A_TILES_K; + coopMatStore( + temp_A[s], Ash_int8, + tk * A_SLAB_U32 + (tm * MMA_M) * A_STRIDE_U32, + A_STRIDE_U32, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#ifdef WEIGHT_INT4 + if (b_owns) { + // ONE b128 store replaces the 4 scalar ds_write_b32 this column-slab + // previously took (spread across 4 threads). + Bsh_v4[ b_v4_off] = uvec4( + widen_nibbles(uint(temp_B[0][b_comp1]), b_par1), + widen_nibbles(uint(temp_B[1][b_comp1]), b_par1), + widen_nibbles(uint(temp_B[2][b_comp1]), b_par1), + widen_nibbles(uint(temp_B[3][b_comp1]), b_par1)); + } +#else + if (b_active) { + const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); + const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); + const uint n_col_base = b_n_uint_col * 4u; + [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { + // bv4: identical address, expressed as (uvec4 element, component). + // This INT8 path writes 4 DIFFERENT columns for one k4, so it cannot + // use a wide store -- it keeps scalar component writes. + Bsh_v4[slab_idx * WG_TILE_N + (n_col_base + n_in_blk)][k4_in_slab] = + uint(temp_B[n_in_blk]); + } + } +#endif + } + + // ========================================================= + // MAIN LOOP (dbuf4) — nested groups x chunks (kept nested; flattening it + // with a conditional coopmat epilog crashes the Xclipse PAL compiler at + // large spec-resolved trip counts). One barrier per chunk. Chunk + // iteration (global index `chunk`): + // 1. barrier — A/B slice (chunk%2) fully written; on the first chunk + // of group g, wsc slice (g%2) is too. + // 2. prefetch — chunk+1 (A tiles, B blocks) into temp; when chunk+1 + // starts a new group, also its wsc element. Skipped + // entirely on the final chunk. + // 3. int8 MMA — on slice (chunk%2) into accum_int32. + // 4. store — temp -> A/B slice ((chunk+1)%2), unpacking the weight; + // on a group boundary, wsc -> slice ((g+1)%2). + // The group epilog runs unconditionally at the tail of each group. + // ========================================================= + uint chunk = 0; + for (uint group_i = 0; group_i < num_groups; ++group_i) { + for (uint inner = 0; inner < CHUNKS_PER_GROUP; ++inner, ++chunk) { + const bool has_next = chunk + 1u < num_chunks; + const bool group_crossing = has_next && (inner + 1u == CHUNKS_PER_GROUP); + const uint cur_a = (chunk % 2u) * ASH_SLICE_U32; + const uint cur_b = (chunk % 2u) * BSH_SLICE_U32; + const uint cur_b_v4 = (chunk % 2u) * BSH_SLICE_V4; + const uint nxt_a = ((chunk + 1u) % 2u) * ASH_SLICE_U32; + const uint nxt_b = ((chunk + 1u) % 2u) * BSH_SLICE_V4; + + // coopmat-lds-fence: barrier() alone does NOT order shared stores against a + // subsequent coopMatLoad on the M51 Xclipse/AMD-PAL driver -- symptom is one + // stale MMA_M-row band of A, all columns, ~2.5% of runs, no crash (observed + // 2026-09-02 in sdpa_compute_out_coopmat.glsl). Measured cost of the fence: + // none (see this change's results). See memory + // `coopmat-lds-needs-explicit-memorybarriershared`. + memoryBarrierShared(); + barrier(); + + // --- 2. prefetch chunk+1 -> temp --- + if (has_next) { + const uint chunkK_nxt = (chunk + 1u) * WG_TILE_K; + // A staging (dbuf4tr's technique): coopMatLoad straight from global. + [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { + const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; + if (t < NUM_A_TILES) { + const uint tm = t / A_TILES_K; + const uint tk = t % A_TILES_K; + coopMatLoad( + temp_A[s], t_packed_int8_input, + (tile_m_start + tm * MMA_M) * a_row_stride_i8 + chunkK_nxt + + tk * MMA_K, + a_row_stride_i8, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#ifdef WEIGHT_INT4 + if (b_owns) { + [[unroll]] for (uint kk = 0; kk < (MMA_K >> 2u); ++kk) { + const uint k4_blk = (chunkK_nxt >> 2u) + b_k4base + kk; +#ifdef WEIGHT_BUFFER + temp_B[kk] = t_packed_weight[(b_n8blk1 * nblocks_x_A) + k4_blk]; +#else + temp_B[kk] = texelFetch(t_packed_weight, ivec2(k4_blk, b_n8blk1), 0); +#endif + } + } + if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { + const uint n_idx = tile_n_start + gl_LocalInvocationID.x; + f16vec4 sv = t_weight_scales[(group_i + 1u) * N4 + (n_idx >> 2u)]; + temp_wsc = float(sv[n_idx & 3u]); + } +#else + if (b_active) { + const uint block_y_w = (chunkK_nxt >> 2u) + b_k4_in_chunk; + const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; +#ifdef WEIGHT_BUFFER + temp_B = t_packed_weight[(block_y_w * N4) + block_x_w]; +#else + temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, block_y_w), 0); +#endif + } +#endif + } + + // --- 3. int8 MMA on the cur slice --- + [[unroll]] for (uint k = 0; k < NUM_K_SLABS; ++k) { + const uint slab_a_base_u32 = cur_a + k * A_SLAB_U32; + const uint slab_b_base_u32 = cur_b + k * B_SLAB_U32; + + coopmat matA[MMAS_PER_SG_M]; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + const uint row_a = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + coopMatLoad( + matA[i], Ash_int8, + slab_a_base_u32 + row_a * A_STRIDE_U32, + A_STRIDE_U32, + gl_CooperativeMatrixLayoutRowMajor); + } + + coopmat matB; + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint col_b = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopMatLoad( + matB, Bsh_v4, + cur_b_v4 + k * WG_TILE_N + col_b, + 1u, + gl_CooperativeMatrixLayoutColumnMajor); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + accum_int32[i][j] = coopMatMulAdd(matA[i], matB, accum_int32[i][j]); + } + } + } + + // --- 4. store temp (chunk+1) -> nxt slice --- + if (has_next) { + // A staging (dbuf4tr's technique): coopMatStore into the nxt slice. + [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { + const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; + if (t < NUM_A_TILES) { + const uint tm = t / A_TILES_K; + const uint tk = t % A_TILES_K; + coopMatStore( + temp_A[s], Ash_int8, + nxt_a + tk * A_SLAB_U32 + (tm * MMA_M) * A_STRIDE_U32, + A_STRIDE_U32, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#ifdef WEIGHT_INT4 + if (b_owns) { + // ONE b128 store replaces the 4 scalar ds_write_b32 this column-slab + // previously took (spread across 4 threads). + Bsh_v4[nxt_b + b_v4_off] = uvec4( + widen_nibbles(uint(temp_B[0][b_comp1]), b_par1), + widen_nibbles(uint(temp_B[1][b_comp1]), b_par1), + widen_nibbles(uint(temp_B[2][b_comp1]), b_par1), + widen_nibbles(uint(temp_B[3][b_comp1]), b_par1)); + } + if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { + const uint wbase_nxt = ((group_i + 1u) % 2u) * WG_TILE_N; + wsc_sh[wbase_nxt + gl_LocalInvocationID.x] = temp_wsc; + } +#else + if (b_active) { + const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); + const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); + const uint n_col_base = b_n_uint_col * 4u; + [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { + Bsh_v4[nxt_b + slab_idx * WG_TILE_N + (n_col_base + n_in_blk)][k4_in_slab] = + uint(temp_B[n_in_blk]); + } + } +#endif + } + } // chunks + + // --- Group epilog: scale-only accumulate, reset accum --- + // dbuf4zpg's, unchanged. Just result += float(acc) * wsc. The + // zero-point subtract and the ifs multiply are hoisted out of the group + // loop (applied once below). + { + const uint wbase = (group_i % 2u) * WG_TILE_N; + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + + coopmat wsc_bcast; + coopMatLoad( + wsc_bcast, wsc_sh, + wbase + local_n_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutRowMajor); + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + result[i][j] += + coopmat( + accum_int32[i][j]) * wsc_bcast; + accum_int32[i][j] = coopmat(0); + } + } + } + } // groups + + // --- Hoisted correction, applied ONCE: -------------------------------- + // result = ifs * ( result - izp * SUM_g wsc_g*wsum_g ) + // dbuf4zpg's, unchanged. izp/ifs are loaded here rather than before the + // group loop so they are not live across it. + { + coopmat + izpf_bcast[MMAS_PER_SG_M]; + coopmat + ifs_bcast[MMAS_PER_SG_M]; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + const uint local_m_base = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + coopmat izp_i; + coopMatLoad( + izp_i, izp_sh, local_m_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutColumnMajor); + izpf_bcast[i] = + coopmat(izp_i); + coopMatLoad( + ifs_bcast[i], ifs_sh, local_m_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutColumnMajor); + } + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopmat wcorr_bcast; + coopMatLoad( + wcorr_bcast, wcorr_sh, local_n_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutRowMajor); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + result[i][j] = ifs_bcast[i] * (result[i][j] - izpf_bcast[i] * wcorr_bcast); + } + } + } + + // --- Bias (optional) --- +#ifdef HAS_BIAS + if (apply_bias > 0) { + for (uint t = gl_LocalInvocationID.x; t < WG_TILE_N; t += WG_SIZE) { + bias_sh[t] = float(t_bias[tile_n_start + t]); + } + memoryBarrierShared(); + barrier(); + } +#endif + + // --- Store result tile --- + // N for the store address math MUST come from the spec constant, not the + // sizes UBO (see out_N_arg above). +#ifdef IO_TEXTURE + // Epilogue iteration i drains accumulator row-block i from EVERY subgroup + // into Csh_out at once, so the SG_GRID_Y bands it holds are disjoint global + // row ranges; the whole workgroup then imageStores them. lr / MMA_M is the + // writing subgroup's warpInTile.y, so the global row reproduces the buffer + // path's gi exactly. + // + // PORTABILITY NOTE: the barrier() in the loop body keeps this loop rolled + // despite [[unroll]], so result[i][j] IS dynamically indexed. Coopmat arrays + // are opaque per-lane storage and dynamic indexing is exactly the construct + // the Xclipse/AMD-PAL compiler has broken before -- check this first if the + // texture variants miscompile on M51. + const uint CSH_TEXELS_PER_ROW = WG_TILE_N / 4u; + const uint CSH_TEXELS = CSH_ROWS * CSH_TEXELS_PER_ROW; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + // Guards Csh_out against the previous iteration's readers. Inert on i == 0 + // but must stay unconditional to remain workgroup-uniform. + // coopmat-lds-fence: barrier() alone does NOT order shared stores against a + // subsequent coopMatLoad on the M51 Xclipse/AMD-PAL driver -- symptom is one + // stale MMA_M-row band of A, all columns, ~2.5% of runs, no crash (observed + // 2026-09-02 in sdpa_compute_out_coopmat.glsl). Measured cost of the fence: + // none (see this change's results). See memory + // `coopmat-lds-needs-explicit-memorybarriershared`. + memoryBarrierShared(); + barrier(); + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { +#ifdef HAS_BIAS + if (apply_bias > 0) { + const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopmat bias_tile; + coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); + result[i][j] += bias_tile; + } +#endif + coopmat out_tile = + coopmat(result[i][j]); + coopMatStore( + out_tile, Csh_out, + warpInTile.y * MMA_M * WG_TILE_N + + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j), + WG_TILE_N, + gl_CooperativeMatrixLayoutRowMajor); + } + memoryBarrierShared(); + barrier(); + + for (uint t = gl_LocalInvocationID.x; t < CSH_TEXELS; t += WG_SIZE) { + const uint lr = t / CSH_TEXELS_PER_ROW; + const uint lc4 = t % CSH_TEXELS_PER_ROW; + const uint m = + tile_m_start + (lr / MMA_M) * SG_TILE_M + i * MMA_M + (lr % MMA_M); + const uint base = lr * WG_TILE_N + lc4 * 4u; + imageStore( + t_output, + ivec3(tile_n_start / 4u + lc4, m, 0), + vec4( + float(Csh_out[base]), + float(Csh_out[base + 1u]), + float(Csh_out[base + 2u]), + float(Csh_out[base + 3u]))); + } + } +#else + const uint N_out = uint(out_N_arg); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint gi = tile_m_start + MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + const uint gj = tile_n_start + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + +#ifdef HAS_BIAS + if (apply_bias > 0) { + const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopmat bias_tile; + coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); + result[i][j] += bias_tile; + } +#endif + + coopmat out_tile = + coopmat(result[i][j]); + coopMatStore( + out_tile, t_output, + gi * N_out + gj, N_out, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#endif // IO_TEXTURE +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbv4.yaml b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbv4.yaml new file mode 100644 index 00000000000..c1ca672d347 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbv4.yaml @@ -0,0 +1,68 @@ +# !!! KNOWN-INCORRECT -- DO NOT PROMOTE, DO NOT TIME !!! +# Fails the correctness gate deterministically (3/3 reps) on exactly the +# num_groups == 2 shapes (K=256: M128_K256_N128, M256_K256_N256). Every K=128, +# K=2048 and K=4096 case PASSES. Root cause NOT yet identified: the +# (slab, col) -> (n8blk, k4, component, parity) mapping was re-derived and is +# provably identical to the baseline slot map, so the defect is in something +# adjacent to the ownership change, not in the address arithmetic. +# Kept opt-in only (never a default) as a recorded negative result. +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# "zpg" + "-tr" combination: tsweep_dbuf4zpg with its per-thread scalar +# A-staging replaced by tsweep_dbuf4tr's coopMat-mediated A staging. B +# staging, zp-hoist, byte-parallel nibble widening are all dbuf4zpg's, +# unchanged -- only A-staging differs. Requires t_packed_int8_input in the +# ROW-MAJOR kPackedInt8_4W layout (same requirement as tsweep_dbuf4tr). +# +# PROMOTED 2026-09-01 as the shipped default -- see QuantizedLinear.cpp's +# dq8ca_coopmat_variant() for the full validation record (10/10 buffer, 6/6 +# texture3d correctness across 1B/3B/8B; +4.2% real e2e prefill; 46.50% -> +# 49.94% efficiency on 8B). Also selectable explicitly via +# ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4zpgtr_txkgs. +# See openspec/changes/archive/2026-08-31-dq8ca-tr-staged-a-on-zpg. +# +# Single seed tile: the prior dbuf4zpg default's own tile (t128x64k32g42s32). +# A re-sweep against this shader's own (lower) register-pressure profile was +# run as a follow-up (coopmat-tr-tilesweep-4w-port) and found no better tile +# -- this remains the best known geometry. + +linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbv4: + parameter_names_with_default_values: + PRECISION: highp + HAS_BIAS: false + IO_STORAGE: buffer + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + MMA_M: 16 + MMA_N: 16 + MMA_K: 16 + WG_TILE_M: 64 + WG_TILE_N: 32 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + shader_variants: + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbv4_t128x64k32g42s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbv4_t128x64k32g42s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw.glsl b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw.glsl new file mode 100644 index 00000000000..2331b04734f --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw.glsl @@ -0,0 +1,781 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +/* + * "zpg" + "-tr" combination: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpg.glsl + * with its per-thread scalar A-staging replaced by + * linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr.glsl's coopMat-mediated A + * staging, PLUS "bw": B staging ownership moved from (block,col) slots spread + * over 4 threads to ONE thread per (K-slab, column), writing its 4 uints at + * consecutive LDS indices so they are merge-eligible. Shared-array types, the + * compute-side loads, the quantization math and every safeguard are + * dbuf4zpgtr's, unchanged. This is + * an ADDITIVE combination, not a redesign: every non-A-staging block below + * (B staging: coalesced write, no skew; zp-hoist: izp/ifs applied once after + * the group loop via wcorr_sh; byte-parallel nibble widening; static + * A_ALWAYS_ACTIVE branch elision -- N/A here, see below; group epilog; + * bias/store epilogue) is byte-identical to dbuf4zpg's. Only the A-staging + * block (prologue load+store, main-loop prefetch+store) is dbuf4tr's, + * verbatim. + * + * dbuf4zpg's per-thread A staging used an `a_active` guard (statically always + * true when A_ACTIVE_THREADS == WG_SIZE, via the A_MAP_FULL-gated + * A_ALWAYS_ACTIVE macro). dbuf4tr's per-SUBGROUP tile map has no equivalent + * concept -- every subgroup participates via a `t < NUM_A_TILES` guard that + * depends only on gl_SubgroupID, not gl_LocalInvocationID.x -- so + * A_MAP_FULL/A_ALWAYS_ACTIVE is dropped entirely in this file; it would be + * dead code for the new A-staging block. + * + * Rationale for combining this way (not the reverse) and why this is worth + * building at all: see this change's design.md D0-D3. In short -- the only + * existing measurement of dbuf4tr's A-staging technique (28.72-30.51%, + * dq8ca-arch-redesign) was taken against dbuf4tr's own pre-zpg baseline + * (old B skew, no byte-parallel widening, no branch elision) -- a materially + * weaker shader than the 46.49-46.50% dbuf4zpg this file now combines it + * with. This file exists to answer whether that combination performs + * differently now that register pressure is already reduced. + * + * A staging (the actual delta from dbuf4zpg): + * dbuf4zpg: per-thread (m4, k4) ivec4 fetch, hoisted a_lds_off0/a_glb_row; + * only A_ACTIVE_THREADS invocations participate, each scattering + * 4 rows into Ash_int8 with 4 scalar stores. + * this file: per-SUBGROUP MMA_M x MMA_K tile fetch via coopMatLoad straight + * from a ROW-MAJOR (kPackedInt8_4W) int8 activation buffer, then + * coopMatStore into the same Ash_int8 slot -- dbuf4tr's mapping, + * unmodified (not re-derived; see design.md D3). + * + * t_packed_int8_input is therefore bound the same way dbuf4tr binds it: a + * SCALAR int8_t array in the kPackedInt8_4W layout (plain row-major int8, + * row stride K), produced by quantize_and_pack_4w_with_group_sums.glsl. + * QuantizedLinear.cpp's dq8ca_variant_wants_rowmajor_a() must recognize this + * file's variant token (tsweep_dbuf4zpgtr_t...) the same way it already + * recognizes tsweep_dbuf4tr_t/trm_t/trd_t, so graph-build time (packer + * selection) and dispatch time (kernel selection) cannot disagree. + * + * B CANNOT be coopmat-staged (int4 nibble unpack; a coopmat's per-lane layout + * is opaque to hand-assembly from unpacked registers) -- unchanged from both + * parent files. B staging below is dbuf4zpg's byte-parallel, coalesced, + * no-skew version, untouched. + * + * The loop structure is dbuf4's (both parents share it), unchanged: + * prologue: prefetch chunk 0 -> temp, store to slice 0 (no barrier) + * per iter: barrier -> prefetch(next) -> MMA(cur) -> store(next) + * kept nested (groups x chunks) with an unconditional group epilog -- + * flattening it crashes the Xclipse PAL compiler at large spec-resolved trip + * counts (see dbuf2's own header). + * + * Selected via + * ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4zpgtr_txkgs<32|64> + * (QuantizedLinear.cpp), additive to the tsweep_dbuf4zpg_t..., tsweep_dbuf4tr_t... + * and tsweep_t... namespaces. NOT the default -- unvalidated until it passes + * repeated test_llama_microbench --correctness-only runs (see + * dq8ca_coopmat_variant()'s comment on why a single pass is not proof). + * + * Performs: out[M,N] = dequant(int8_act) * dequant(int_w) (+ bias) + * via coopmat x coopmat -> coopmat on the matrix unit. + * + * Hard preconditions (dbuf4zpg's, plus dbuf4tr's row-major/alignment ones): + * M % WG_TILE_M == 0, N % WG_TILE_N == 0, K % WG_TILE_K == 0, + * group_size % WG_TILE_K == 0, K % 4 == 0, + * WG_TILE_M % MMA_M == 0, WG_TILE_K % MMA_K == 0, + * t_packed_int8_input in kPackedInt8_4W (row-major) layout, + * device exposes coopmatx-> at 16x16x16. + */ + +#version 450 core + +#extension GL_KHR_cooperative_matrix : require +#extension GL_KHR_memory_scope_semantics : require +#extension GL_KHR_shader_subgroup_basic : enable +#extension GL_EXT_shader_explicit_arithmetic_types : require +#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require +// 8-bit SSBO access: A is bound as a scalar int8_t array so that the +// coopMatLoad below has a MATCHING component type (see dbuf4tr's header for +// why the type must match on this driver). +#extension GL_EXT_shader_8bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_EXT_control_flow_attributes : enable + +#define PRECISION ${PRECISION} + +$if WEIGHT_NBITS == 4: + #define WEIGHT_INT4 + +$if HAS_BIAS: + #define HAS_BIAS + +$if WEIGHT_STORAGE == "buffer": + #define WEIGHT_BUFFER + +$if IO_STORAGE == "texture3d": + #define IO_TEXTURE + +layout(std430) buffer; + +#include "common.glslh" + +// Bindings — match add_linear_dqa_qw_node arg order: +// output(0), fp_input(1), packed_int8_input(2), int_input_sums(3 - unused), +// input_scales(4), input_zps(5), packed_weight(6), weight_sums(7), +// weight_scales(8), bias(9). +${layout_declare_tensor(B, "w", "t_output", "half", IO_STORAGE, is_scalar_array=True)} +// t_input is unread here -- the activations arrive already quantized in +// t_packed_int8_input -- but stays declared so the binding layout matches the +// dispatch site. It tracks IO_STORAGE so the two IO tensors stay consistent. +${layout_declare_tensor(B, "r", "t_input", "half", IO_STORAGE, is_scalar_array=False)} +// ROW-MAJOR (kPackedInt8_4W) packed activations, bound as a scalar int8_t +// array (row stride = K int8) -- dbuf4tr's binding, unchanged. The stock +// 4h4w layout dbuf4zpg uses is NOT row-major (component index selects a row, +// non-affine), so it cannot be addressed by any coopMatLoad. +${layout_declare_tensor(B, "r", "t_packed_int8_input", "int8", "buffer", is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_int8_input_sums", "int", "buffer", is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_int8_input_scales", "half", "texture3d")} +${layout_declare_tensor(B, "r", "t_int8_input_zps", "int8", "texture3d")} +${layout_declare_tensor(B, "r", "t_packed_weight", "int", WEIGHT_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_weight_sums", "int", "buffer", is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_weight_scales", "half", "buffer", is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_bias", "half", "buffer", is_scalar_array=True)} + +${layout_declare_ubo(B, "ivec4", "output_sizes")} +${layout_declare_ubo(B, "ivec4", "input_sizes")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +${layout_declare_spec_const(C, "int", "apply_bias", "0")} +// INT4 only; inert (0) for INT8 so the dispatcher's spec list lines up. +${layout_declare_spec_const(C, "int", "K4_per_group", "0")} +${layout_declare_spec_const(C, "int", "num_groups_arg", "0")} +${layout_declare_spec_const(C, "int", "out_N_arg", "0")} + +// Tile geometry +const uint MMA_M = ${MMA_M}; +const uint MMA_N = ${MMA_N}; +const uint MMA_K = ${MMA_K}; + +const uint WG_TILE_M = ${WG_TILE_M}; +const uint WG_TILE_N = ${WG_TILE_N}; +const uint WG_TILE_K = ${WG_TILE_K}; + +const uint SG_GRID_X = ${SG_GRID_X}; +const uint SG_GRID_Y = ${SG_GRID_Y}; +const uint SUBGROUP_SIZE = ${SUBGROUP_SIZE}; +const uint NUM_SUBGROUPS = SG_GRID_X * SG_GRID_Y; +const uint WG_SIZE = NUM_SUBGROUPS * SUBGROUP_SIZE; + +const uint SG_TILE_M = WG_TILE_M / SG_GRID_Y; +const uint SG_TILE_N = WG_TILE_N / SG_GRID_X; +const uint MMAS_PER_SG_M = SG_TILE_M / MMA_M; +const uint MMAS_PER_SG_N = SG_TILE_N / MMA_N; + +const uint A_SLAB_INT8 = WG_TILE_M * MMA_K; +const uint B_USEFUL_U32 = MMA_K / 4u; +// No skew + coalesced write -- dbuf4zpg's B fix, unchanged (this file does +// not touch B staging at all). +const uint B_STRIDE_U32 = B_USEFUL_U32; +const uint B_SLAB_U32 = WG_TILE_N * B_STRIDE_U32; +const uint NUM_K_SLABS = WG_TILE_K / MMA_K; + +const uint A_SLAB_U32 = A_SLAB_INT8 / 4u; +const uint A_STRIDE_U32 = MMA_K / 4u; + +// One ping-pong slice covers all K-slabs of one chunk. +const uint ASH_SLICE_U32 = NUM_K_SLABS * A_SLAB_U32; +const uint BSH_SLICE_U32 = NUM_K_SLABS * B_SLAB_U32; + +// Double-buffered MMA operand staging. +shared uint Ash_int8[2u * ASH_SLICE_U32]; +shared uint Bsh_int8[2u * BSH_SLICE_U32]; + +// Per-WG-tile-row activation params (loaded ONCE at WG start; constant +// across groups). +shared int izp_sh[WG_TILE_M]; // int32 (cast from int8 source) for broadcast +shared float ifs_sh[WG_TILE_M]; // float32 (cast from fp16 source) for broadcast + +// Per-(group, output-channel) weight params, ping-ponged by group parity. +// (For per-channel INT8 only slice 0 is ever used.) +shared float wsc_sh[2u * WG_TILE_N]; +// SUM_g wsc[g][n]*wsum[g][n] per output channel -- weight-side only, so it is +// accumulated once in the prologue. dbuf4zpg's zp-hoist, unchanged. +shared float wcorr_sh[WG_TILE_N]; + +#ifdef HAS_BIAS +shared float bias_sh[WG_TILE_N]; +#endif + +#ifdef IO_TEXTURE +// Result staging for the imageStore epilogue, mirroring the fp16 kernel: +// SG_GRID_Y bands of MMA_M rows, each WG_TILE_N wide, row-major. A full +// WG_TILE_M x WG_TILE_N buffer would cost SG_GRID_Y/MMAS_PER_SG_M x more LDS +// and wreck occupancy. float16_t-typed because coopMatStore needs it. +const uint CSH_ROWS = SG_GRID_Y * MMA_M; +shared float16_t Csh_out[CSH_ROWS * WG_TILE_N]; +#endif + +// Running fp32 accumulator (across all groups). +coopmat + result[MMAS_PER_SG_M][MMAS_PER_SG_N]; + +// Per-group int32 MMA accumulator. +coopmat + accum_int32[MMAS_PER_SG_M][MMAS_PER_SG_N]; + + +// Byte-parallel int4 -> int8 widening. dbuf4zpg's, unchanged (B-side only). +// +// The four nibbles this shader needs from one packed uint are ALREADY one per +// byte (bits 3:0 of each byte for parity 0, bits 7:4 for parity 1), so all four +// can be widened at once instead of with a per-nibble +// shift/mask/bias-subtract/mask chain. +// +// For v in [0,15] the biased value is v-8. `v ^ 8` is exactly the 4-bit +// two's-complement pattern of v-8, because -8 == +8 (mod 16): +// v=0 -> 0x8 -> -8 v=7 -> 0xF -> -1 +// v=8 -> 0x0 -> 0 v=15 -> 0x7 -> +7 +// so the only remaining work is sign-extending bit 3 into bits 7:4 per byte. +// `sgn * 0x1E` does that with no cross-byte carry: 0x08 * 0x1E == 0xF0 exactly, +// and sgn is at most 0x08080808 so the product is at most 0xF0F0F0F0. +// +// A naive `nib - 0x08080808` would NOT work -- it borrows across byte lanes +// whenever a nibble is < 8. Shifts must be on uint, not int, so the >> is +// logical rather than arithmetic. +// +// ~5 ops per 4 weights vs ~22 for the per-nibble chain; bit-identical output. +uint widen_nibbles(const uint w, const uint parity) { + const uint nib = (parity == 0u) ? (w & 0x0F0F0F0Fu) : ((w >> 4u) & 0x0F0F0F0Fu); + const uint p = nib ^ 0x08080808u; + const uint sgn = p & 0x08080808u; + return p | (sgn * 0x1Eu); +} + +void main() { + const uvec2 tileID = uvec2(gl_WorkGroupID.xy); + const uvec2 warpInTile = uvec2( + gl_SubgroupID % SG_GRID_X, + gl_SubgroupID / SG_GRID_X); + + const uint K = uint(input_sizes.x); + const uint N = uint(output_sizes.x); + const uint N4 = (N + 3u) / 4u; + const uint nblocks_x_A = (K + 3u) >> 2u; + // A row stride in INT8 elements (dbuf4tr's binding is row-major int8, not + // the 4h4w ivec4 block layout dbuf4zpg used -- so A addressing below is in + // int8 elements, not int, and derived from nblocks_x_A so it matches the + // packer's `m_row * K4 + k4` addressing exactly (K % 4 == 0 makes them + // equal to K directly). + const uint a_row_stride_i8 = nblocks_x_A * 4u; + +#ifdef WEIGHT_INT4 + const uint num_groups = uint(num_groups_arg); + const uint CHUNKS_PER_GROUP = uint(K4_per_group) * 4u / WG_TILE_K; +#else + // Per-channel: a single quant "group" spanning all of K. The nested + // groups x chunks loop below collapses to a flat chunk loop, the wsum/wsc + // ping-pong never crosses a boundary, and the epilog runs exactly once. + const uint num_groups = 1u; + const uint CHUNKS_PER_GROUP = uint(num_groups_arg); +#endif + const uint num_chunks = num_groups * CHUNKS_PER_GROUP; + + const uint tile_m_start = WG_TILE_M * tileID.y; + const uint tile_n_start = WG_TILE_N * tileID.x; + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + result[i][j] = coopmat(0.0); + accum_int32[i][j] = coopmat(0); + } + } + + const uint K_BLOCKS_PER_CHUNK = WG_TILE_K >> 2u; + + // --- A staging tile map (dbuf4tr's, unmodified): one MMA_M x MMA_K coopmat + // tile per subgroup per slot, dealt round-robin across the + // NUM_SUBGROUPS subgroups so every subgroup participates. Replaces + // dbuf4zpg's per-thread (m4, k4) map / a_active guard entirely -- see + // design.md D3 for why this is reused as-is, not re-derived. + const uint A_TILES_M = WG_TILE_M / MMA_M; + const uint A_TILES_K = WG_TILE_K / MMA_K; // == NUM_K_SLABS + const uint NUM_A_TILES = A_TILES_M * A_TILES_K; + const uint A_TILES_PER_SG = (NUM_A_TILES + NUM_SUBGROUPS - 1u) / NUM_SUBGROUPS; + +#ifdef WEIGHT_INT4 + // --- B staging thread map: (block, col) slots; each slot extracts one + // ColumnMajor LDS uint (4 K-contiguous sign-extended int8) --- + const uint B_TOTAL_SLOTS = K_BLOCKS_PER_CHUNK * WG_TILE_N; + const uint B_SLOTS_PER_THREAD = B_TOTAL_SLOTS / WG_SIZE; + const uint N8_PER_TILE = WG_TILE_N >> 3u; +#else + // --- B staging thread map: one (k4, n4) ivec4 block per active thread --- + const uint B_FETCH_SLOTS = K_BLOCKS_PER_CHUNK * (WG_TILE_N >> 2u); + const uint N4_PER_TILE = WG_TILE_N >> 2u; + const uint b_k4_in_chunk = gl_LocalInvocationID.x / N4_PER_TILE; + const uint b_n_uint_col = gl_LocalInvocationID.x % N4_PER_TILE; + const bool b_active = gl_LocalInvocationID.x < B_FETCH_SLOTS; +#endif + + // ===== INTERVENTION F: hoist loop-invariant B staging index math ===== + // dbuf4zpg's, unchanged -- B staging is untouched by this file's A-staging + // swap. See dbuf4zpg's header for the full rationale (ablation-attributed + // -16.8% block, mostly index arithmetic not memory traffic). +#ifdef WEIGHT_INT4 + // ===== bw: SOURCE-side change only -- one thread owns a whole (slab, col) + // column-slab and writes its 4 uints at CONSECUTIVE indices, so the four + // ds_write_b32 are merge-eligible. The shared array stays `uint` (the proven + // 4x packing ratio) and the compute-side coopMatLoad is byte-for-byte + // unchanged -- deliberately NOT the uvec4 retype, which fails correctness on + // this driver (16x packing; see the 8x coopMatStore precedent). + // comp/par depend only on the column, so a thread's 4 fetches share them and + // differ only in k4: 4 consecutive k4 blocks of the same n8 texel column. + const uint B_OWNERS = NUM_K_SLABS * WG_TILE_N; + const bool b_owns = gl_LocalInvocationID.x < B_OWNERS; + const uint b_slab = gl_LocalInvocationID.x / WG_TILE_N; + const uint b_col = gl_LocalInvocationID.x % WG_TILE_N; + const uint b_rem = b_col & 7u; + const uint b_comp1 = b_rem & 3u; + const uint b_par1 = b_rem >> 2u; + const uint b_n8blk1 = (tile_n_start >> 3u) + (b_col >> 3u); + const uint b_k4base = b_slab * (MMA_K >> 2u); + // same address the destination-oriented map produced, just all 4 in one thread + const uint b_lds_base = b_slab * B_SLAB_U32 + b_col * B_STRIDE_U32; +#endif + + // Prefetch temp registers. temp_A is a coopmat array (dbuf4tr's A-staging + // technique); indices into it are [[unroll]]-resolved compile-time + // constants, never dynamic -- dynamic indexing of a coopmat array is + // exactly the construct the Xclipse/AMD-PAL compiler has miscompiled + // before. + coopmat + temp_A[A_TILES_PER_SG]; +#ifdef WEIGHT_INT4 + ivec4 temp_B[MMA_K >> 2u]; // bw: one per k4 of the owned column-slab + float temp_wsc; +#else + ivec4 temp_B; +#endif + + // ========================================================= + // PROLOGUE + // ========================================================= + if (gl_LocalInvocationID.x < (WG_TILE_M >> 2u)) { + const uint m4 = (tile_m_start >> 2u) + gl_LocalInvocationID.x; + const vec4 sc = vec4(texelFetch(t_int8_input_scales, ivec3(m4, 0, 0), 0)); + const ivec4 zp = texelFetch(t_int8_input_zps, ivec3(m4, 0, 0), 0); + const uint base = gl_LocalInvocationID.x * 4u; + ifs_sh[base + 0u] = sc.x; ifs_sh[base + 1u] = sc.y; + ifs_sh[base + 2u] = sc.z; ifs_sh[base + 3u] = sc.w; + izp_sh[base + 0u] = zp.x; izp_sh[base + 1u] = zp.y; + izp_sh[base + 2u] = zp.z; izp_sh[base + 3u] = zp.w; + } + // Group 0 weight scales -> slice 0, and the hoisted weight-side correction + // SUM_g wsc[g][n]*wsum[g][n] accumulated across ALL groups. dbuf4zpg's + // zp-hoist, unchanged. + if (gl_LocalInvocationID.x < WG_TILE_N) { + const uint n_idx = tile_n_start + gl_LocalInvocationID.x; + f16vec4 sv0 = t_weight_scales[n_idx >> 2u]; + wsc_sh[gl_LocalInvocationID.x] = float(sv0[n_idx & 3u]); + + float corr = 0.0; + for (uint g = 0; g < num_groups; ++g) { + f16vec4 sv = t_weight_scales[g * N4 + (n_idx >> 2u)]; + corr += float(sv[n_idx & 3u]) * float(t_weight_sums[g * N + n_idx]); + } + wcorr_sh[gl_LocalInvocationID.x] = corr; + } + memoryBarrierShared(); + barrier(); + + // NOTE: dbuf4zpg builds izp_bcast/ifs_bcast AFTER the group loop, not here + // -- that is the register-pressure saving zp-hoist buys. Unchanged. + + // dbuf4: prefetch chunk 0 into temp registers, THEN store to slice 0 (no + // barrier here -- the main loop's first iteration barriers before + // reading slice 0). + // + // A staging (dbuf4tr's technique): per-subgroup coopMatLoad straight from + // the row-major global buffer. + [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { + const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; + if (t < NUM_A_TILES) { + const uint tm = t / A_TILES_K; + const uint tk = t % A_TILES_K; + coopMatLoad( + temp_A[s], t_packed_int8_input, + (tile_m_start + tm * MMA_M) * a_row_stride_i8 + tk * MMA_K, + a_row_stride_i8, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#ifdef WEIGHT_INT4 + if (b_owns) { + [[unroll]] for (uint kk = 0; kk < (MMA_K >> 2u); ++kk) { +#ifdef WEIGHT_BUFFER + temp_B[kk] = t_packed_weight[(b_n8blk1 * nblocks_x_A) + b_k4base + kk]; +#else + temp_B[kk] = texelFetch(t_packed_weight, ivec2(b_k4base + kk, b_n8blk1), 0); +#endif + } + } +#else + if (b_active) { + const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; +#ifdef WEIGHT_BUFFER + temp_B = t_packed_weight[(b_k4_in_chunk * N4) + block_x_w]; +#else + temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, b_k4_in_chunk), 0); +#endif + } +#endif + { + // store chunk 0 -> slice 0 + // A staging (dbuf4tr's technique): coopMatStore into the same Ash_int8 + // slot layout dbuf4zpg's scalar scatter used to write. + [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { + const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; + if (t < NUM_A_TILES) { + const uint tm = t / A_TILES_K; + const uint tk = t % A_TILES_K; + coopMatStore( + temp_A[s], Ash_int8, + tk * A_SLAB_U32 + (tm * MMA_M) * A_STRIDE_U32, + A_STRIDE_U32, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#ifdef WEIGHT_INT4 + if (b_owns) { + // 4 CONSECUTIVE uints from one thread -> merge-eligible into a wider + // store, versus 4 scalar stores from 4 different threads before. + [[unroll]] for (uint kk = 0; kk < (MMA_K >> 2u); ++kk) { + Bsh_int8[ b_lds_base + kk] = + widen_nibbles(uint(temp_B[kk][b_comp1]), b_par1); + } + } +#else + if (b_active) { + const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); + const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); + const uint n_col_base = b_n_uint_col * 4u; + [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { + Bsh_int8[slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab] = + uint(temp_B[n_in_blk]); + } + } +#endif + } + + // ========================================================= + // MAIN LOOP (dbuf4) — nested groups x chunks (kept nested; flattening it + // with a conditional coopmat epilog crashes the Xclipse PAL compiler at + // large spec-resolved trip counts). One barrier per chunk. Chunk + // iteration (global index `chunk`): + // 1. barrier — A/B slice (chunk%2) fully written; on the first chunk + // of group g, wsc slice (g%2) is too. + // 2. prefetch — chunk+1 (A tiles, B blocks) into temp; when chunk+1 + // starts a new group, also its wsc element. Skipped + // entirely on the final chunk. + // 3. int8 MMA — on slice (chunk%2) into accum_int32. + // 4. store — temp -> A/B slice ((chunk+1)%2), unpacking the weight; + // on a group boundary, wsc -> slice ((g+1)%2). + // The group epilog runs unconditionally at the tail of each group. + // ========================================================= + uint chunk = 0; + for (uint group_i = 0; group_i < num_groups; ++group_i) { + for (uint inner = 0; inner < CHUNKS_PER_GROUP; ++inner, ++chunk) { + const bool has_next = chunk + 1u < num_chunks; + const bool group_crossing = has_next && (inner + 1u == CHUNKS_PER_GROUP); + const uint cur_a = (chunk % 2u) * ASH_SLICE_U32; + const uint cur_b = (chunk % 2u) * BSH_SLICE_U32; + const uint nxt_a = ((chunk + 1u) % 2u) * ASH_SLICE_U32; + const uint nxt_b = ((chunk + 1u) % 2u) * BSH_SLICE_U32; + + // coopmat-lds-fence: barrier() alone does NOT order shared stores against a + // subsequent coopMatLoad on the M51 Xclipse/AMD-PAL driver -- symptom is one + // stale MMA_M-row band of A, all columns, ~2.5% of runs, no crash (observed + // 2026-09-02 in sdpa_compute_out_coopmat.glsl). Measured cost of the fence: + // none (see this change's results). See memory + // `coopmat-lds-needs-explicit-memorybarriershared`. + memoryBarrierShared(); + barrier(); + + // --- 2. prefetch chunk+1 -> temp --- + if (has_next) { + const uint chunkK_nxt = (chunk + 1u) * WG_TILE_K; + // A staging (dbuf4tr's technique): coopMatLoad straight from global. + [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { + const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; + if (t < NUM_A_TILES) { + const uint tm = t / A_TILES_K; + const uint tk = t % A_TILES_K; + coopMatLoad( + temp_A[s], t_packed_int8_input, + (tile_m_start + tm * MMA_M) * a_row_stride_i8 + chunkK_nxt + + tk * MMA_K, + a_row_stride_i8, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#ifdef WEIGHT_INT4 + if (b_owns) { + [[unroll]] for (uint kk = 0; kk < (MMA_K >> 2u); ++kk) { + const uint k4_blk = (chunkK_nxt >> 2u) + b_k4base + kk; +#ifdef WEIGHT_BUFFER + temp_B[kk] = t_packed_weight[(b_n8blk1 * nblocks_x_A) + k4_blk]; +#else + temp_B[kk] = texelFetch(t_packed_weight, ivec2(k4_blk, b_n8blk1), 0); +#endif + } + } + if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { + const uint n_idx = tile_n_start + gl_LocalInvocationID.x; + f16vec4 sv = t_weight_scales[(group_i + 1u) * N4 + (n_idx >> 2u)]; + temp_wsc = float(sv[n_idx & 3u]); + } +#else + if (b_active) { + const uint block_y_w = (chunkK_nxt >> 2u) + b_k4_in_chunk; + const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; +#ifdef WEIGHT_BUFFER + temp_B = t_packed_weight[(block_y_w * N4) + block_x_w]; +#else + temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, block_y_w), 0); +#endif + } +#endif + } + + // --- 3. int8 MMA on the cur slice --- + [[unroll]] for (uint k = 0; k < NUM_K_SLABS; ++k) { + const uint slab_a_base_u32 = cur_a + k * A_SLAB_U32; + const uint slab_b_base_u32 = cur_b + k * B_SLAB_U32; + + coopmat matA[MMAS_PER_SG_M]; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + const uint row_a = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + coopMatLoad( + matA[i], Ash_int8, + slab_a_base_u32 + row_a * A_STRIDE_U32, + A_STRIDE_U32, + gl_CooperativeMatrixLayoutRowMajor); + } + + coopmat matB; + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint col_b = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopMatLoad( + matB, Bsh_int8, + slab_b_base_u32 + col_b * B_STRIDE_U32, + B_STRIDE_U32, + gl_CooperativeMatrixLayoutColumnMajor); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + accum_int32[i][j] = coopMatMulAdd(matA[i], matB, accum_int32[i][j]); + } + } + } + + // --- 4. store temp (chunk+1) -> nxt slice --- + if (has_next) { + // A staging (dbuf4tr's technique): coopMatStore into the nxt slice. + [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { + const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; + if (t < NUM_A_TILES) { + const uint tm = t / A_TILES_K; + const uint tk = t % A_TILES_K; + coopMatStore( + temp_A[s], Ash_int8, + nxt_a + tk * A_SLAB_U32 + (tm * MMA_M) * A_STRIDE_U32, + A_STRIDE_U32, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#ifdef WEIGHT_INT4 + if (b_owns) { + // 4 CONSECUTIVE uints from one thread -> merge-eligible into a wider + // store, versus 4 scalar stores from 4 different threads before. + [[unroll]] for (uint kk = 0; kk < (MMA_K >> 2u); ++kk) { + Bsh_int8[nxt_b + b_lds_base + kk] = + widen_nibbles(uint(temp_B[kk][b_comp1]), b_par1); + } + } + if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { + const uint wbase_nxt = ((group_i + 1u) % 2u) * WG_TILE_N; + wsc_sh[wbase_nxt + gl_LocalInvocationID.x] = temp_wsc; + } +#else + if (b_active) { + const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); + const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); + const uint n_col_base = b_n_uint_col * 4u; + [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { + Bsh_int8[nxt_b + slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab] = + uint(temp_B[n_in_blk]); + } + } +#endif + } + } // chunks + + // --- Group epilog: scale-only accumulate, reset accum --- + // dbuf4zpg's, unchanged. Just result += float(acc) * wsc. The + // zero-point subtract and the ifs multiply are hoisted out of the group + // loop (applied once below). + { + const uint wbase = (group_i % 2u) * WG_TILE_N; + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + + coopmat wsc_bcast; + coopMatLoad( + wsc_bcast, wsc_sh, + wbase + local_n_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutRowMajor); + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + result[i][j] += + coopmat( + accum_int32[i][j]) * wsc_bcast; + accum_int32[i][j] = coopmat(0); + } + } + } + } // groups + + // --- Hoisted correction, applied ONCE: -------------------------------- + // result = ifs * ( result - izp * SUM_g wsc_g*wsum_g ) + // dbuf4zpg's, unchanged. izp/ifs are loaded here rather than before the + // group loop so they are not live across it. + { + coopmat + izpf_bcast[MMAS_PER_SG_M]; + coopmat + ifs_bcast[MMAS_PER_SG_M]; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + const uint local_m_base = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + coopmat izp_i; + coopMatLoad( + izp_i, izp_sh, local_m_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutColumnMajor); + izpf_bcast[i] = + coopmat(izp_i); + coopMatLoad( + ifs_bcast[i], ifs_sh, local_m_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutColumnMajor); + } + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopmat wcorr_bcast; + coopMatLoad( + wcorr_bcast, wcorr_sh, local_n_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutRowMajor); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + result[i][j] = ifs_bcast[i] * (result[i][j] - izpf_bcast[i] * wcorr_bcast); + } + } + } + + // --- Bias (optional) --- +#ifdef HAS_BIAS + if (apply_bias > 0) { + for (uint t = gl_LocalInvocationID.x; t < WG_TILE_N; t += WG_SIZE) { + bias_sh[t] = float(t_bias[tile_n_start + t]); + } + memoryBarrierShared(); + barrier(); + } +#endif + + // --- Store result tile --- + // N for the store address math MUST come from the spec constant, not the + // sizes UBO (see out_N_arg above). +#ifdef IO_TEXTURE + // Epilogue iteration i drains accumulator row-block i from EVERY subgroup + // into Csh_out at once, so the SG_GRID_Y bands it holds are disjoint global + // row ranges; the whole workgroup then imageStores them. lr / MMA_M is the + // writing subgroup's warpInTile.y, so the global row reproduces the buffer + // path's gi exactly. + // + // PORTABILITY NOTE: the barrier() in the loop body keeps this loop rolled + // despite [[unroll]], so result[i][j] IS dynamically indexed. Coopmat arrays + // are opaque per-lane storage and dynamic indexing is exactly the construct + // the Xclipse/AMD-PAL compiler has broken before -- check this first if the + // texture variants miscompile on M51. + const uint CSH_TEXELS_PER_ROW = WG_TILE_N / 4u; + const uint CSH_TEXELS = CSH_ROWS * CSH_TEXELS_PER_ROW; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + // Guards Csh_out against the previous iteration's readers. Inert on i == 0 + // but must stay unconditional to remain workgroup-uniform. + // coopmat-lds-fence: barrier() alone does NOT order shared stores against a + // subsequent coopMatLoad on the M51 Xclipse/AMD-PAL driver -- symptom is one + // stale MMA_M-row band of A, all columns, ~2.5% of runs, no crash (observed + // 2026-09-02 in sdpa_compute_out_coopmat.glsl). Measured cost of the fence: + // none (see this change's results). See memory + // `coopmat-lds-needs-explicit-memorybarriershared`. + memoryBarrierShared(); + barrier(); + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { +#ifdef HAS_BIAS + if (apply_bias > 0) { + const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopmat bias_tile; + coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); + result[i][j] += bias_tile; + } +#endif + coopmat out_tile = + coopmat(result[i][j]); + coopMatStore( + out_tile, Csh_out, + warpInTile.y * MMA_M * WG_TILE_N + + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j), + WG_TILE_N, + gl_CooperativeMatrixLayoutRowMajor); + } + memoryBarrierShared(); + barrier(); + + for (uint t = gl_LocalInvocationID.x; t < CSH_TEXELS; t += WG_SIZE) { + const uint lr = t / CSH_TEXELS_PER_ROW; + const uint lc4 = t % CSH_TEXELS_PER_ROW; + const uint m = + tile_m_start + (lr / MMA_M) * SG_TILE_M + i * MMA_M + (lr % MMA_M); + const uint base = lr * WG_TILE_N + lc4 * 4u; + imageStore( + t_output, + ivec3(tile_n_start / 4u + lc4, m, 0), + vec4( + float(Csh_out[base]), + float(Csh_out[base + 1u]), + float(Csh_out[base + 2u]), + float(Csh_out[base + 3u]))); + } + } +#else + const uint N_out = uint(out_N_arg); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint gi = tile_m_start + MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + const uint gj = tile_n_start + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + +#ifdef HAS_BIAS + if (apply_bias > 0) { + const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopmat bias_tile; + coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); + result[i][j] += bias_tile; + } +#endif + + coopmat out_tile = + coopmat(result[i][j]); + coopMatStore( + out_tile, t_output, + gi * N_out + gj, N_out, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#endif // IO_TEXTURE +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw.yaml b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw.yaml new file mode 100644 index 00000000000..ff97b1a2710 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw.yaml @@ -0,0 +1,68 @@ +# !!! KNOWN-INCORRECT -- DO NOT PROMOTE, DO NOT TIME !!! +# Fails the correctness gate deterministically (3/3 reps) on exactly the +# num_groups == 2 shapes (K=256: M128_K256_N128, M256_K256_N256). Every K=128, +# K=2048 and K=4096 case PASSES. Root cause NOT yet identified: the +# (slab, col) -> (n8blk, k4, component, parity) mapping was re-derived and is +# provably identical to the baseline slot map, so the defect is in something +# adjacent to the ownership change, not in the address arithmetic. +# Kept opt-in only (never a default) as a recorded negative result. +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# "zpg" + "-tr" combination: tsweep_dbuf4zpg with its per-thread scalar +# A-staging replaced by tsweep_dbuf4tr's coopMat-mediated A staging. B +# staging, zp-hoist, byte-parallel nibble widening are all dbuf4zpg's, +# unchanged -- only A-staging differs. Requires t_packed_int8_input in the +# ROW-MAJOR kPackedInt8_4W layout (same requirement as tsweep_dbuf4tr). +# +# PROMOTED 2026-09-01 as the shipped default -- see QuantizedLinear.cpp's +# dq8ca_coopmat_variant() for the full validation record (10/10 buffer, 6/6 +# texture3d correctness across 1B/3B/8B; +4.2% real e2e prefill; 46.50% -> +# 49.94% efficiency on 8B). Also selectable explicitly via +# ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4zpgbw_txkgs. +# See openspec/changes/archive/2026-08-31-dq8ca-tr-staged-a-on-zpg. +# +# Single seed tile: the prior dbuf4zpg default's own tile (t128x64k32g42s32). +# A re-sweep against this shader's own (lower) register-pressure profile was +# run as a follow-up (coopmat-tr-tilesweep-4w-port) and found no better tile +# -- this remains the best known geometry. + +linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw: + parameter_names_with_default_values: + PRECISION: highp + HAS_BIAS: false + IO_STORAGE: buffer + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + MMA_M: 16 + MMA_N: 16 + MMA_K: 16 + WG_TILE_M: 64 + WG_TILE_N: 32 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + shader_variants: + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw_t128x64k32g42s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw_t128x64k32g42s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw2.glsl b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw2.glsl new file mode 100644 index 00000000000..38640aca569 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw2.glsl @@ -0,0 +1,762 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +/* + * "zpg" + "-tr" combination: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpg.glsl + * with its per-thread scalar A-staging replaced by + * linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr.glsl's coopMat-mediated A + * staging (coopMatLoad(global) -> coopmat<> -> coopMatStore(LDS)). This is + * an ADDITIVE combination, not a redesign: every non-A-staging block below + * (B staging: coalesced write, no skew; zp-hoist: izp/ifs applied once after + * the group loop via wcorr_sh; byte-parallel nibble widening; static + * A_ALWAYS_ACTIVE branch elision -- N/A here, see below; group epilog; + * bias/store epilogue) is byte-identical to dbuf4zpg's. Only the A-staging + * block (prologue load+store, main-loop prefetch+store) is dbuf4tr's, + * verbatim. + * + * dbuf4zpg's per-thread A staging used an `a_active` guard (statically always + * true when A_ACTIVE_THREADS == WG_SIZE, via the A_MAP_FULL-gated + * A_ALWAYS_ACTIVE macro). dbuf4tr's per-SUBGROUP tile map has no equivalent + * concept -- every subgroup participates via a `t < NUM_A_TILES` guard that + * depends only on gl_SubgroupID, not gl_LocalInvocationID.x -- so + * A_MAP_FULL/A_ALWAYS_ACTIVE is dropped entirely in this file; it would be + * dead code for the new A-staging block. + * + * Rationale for combining this way (not the reverse) and why this is worth + * building at all: see this change's design.md D0-D3. In short -- the only + * existing measurement of dbuf4tr's A-staging technique (28.72-30.51%, + * dq8ca-arch-redesign) was taken against dbuf4tr's own pre-zpg baseline + * (old B skew, no byte-parallel widening, no branch elision) -- a materially + * weaker shader than the 46.49-46.50% dbuf4zpg this file now combines it + * with. This file exists to answer whether that combination performs + * differently now that register pressure is already reduced. + * + * A staging (the actual delta from dbuf4zpg): + * dbuf4zpg: per-thread (m4, k4) ivec4 fetch, hoisted a_lds_off0/a_glb_row; + * only A_ACTIVE_THREADS invocations participate, each scattering + * 4 rows into Ash_int8 with 4 scalar stores. + * this file: per-SUBGROUP MMA_M x MMA_K tile fetch via coopMatLoad straight + * from a ROW-MAJOR (kPackedInt8_4W) int8 activation buffer, then + * coopMatStore into the same Ash_int8 slot -- dbuf4tr's mapping, + * unmodified (not re-derived; see design.md D3). + * + * t_packed_int8_input is therefore bound the same way dbuf4tr binds it: a + * SCALAR int8_t array in the kPackedInt8_4W layout (plain row-major int8, + * row stride K), produced by quantize_and_pack_4w_with_group_sums.glsl. + * QuantizedLinear.cpp's dq8ca_variant_wants_rowmajor_a() must recognize this + * file's variant token (tsweep_dbuf4zpgtr_t...) the same way it already + * recognizes tsweep_dbuf4tr_t/trm_t/trd_t, so graph-build time (packer + * selection) and dispatch time (kernel selection) cannot disagree. + * + * B CANNOT be coopmat-staged (int4 nibble unpack; a coopmat's per-lane layout + * is opaque to hand-assembly from unpacked registers) -- unchanged from both + * parent files. B staging below is dbuf4zpg's byte-parallel, coalesced, + * no-skew version, untouched. + * + * The loop structure is dbuf4's (both parents share it), unchanged: + * prologue: prefetch chunk 0 -> temp, store to slice 0 (no barrier) + * per iter: barrier -> prefetch(next) -> MMA(cur) -> store(next) + * kept nested (groups x chunks) with an unconditional group epilog -- + * flattening it crashes the Xclipse PAL compiler at large spec-resolved trip + * counts (see dbuf2's own header). + * + * Selected via + * ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4zpgtr_txkgs<32|64> + * (QuantizedLinear.cpp), additive to the tsweep_dbuf4zpg_t..., tsweep_dbuf4tr_t... + * and tsweep_t... namespaces. NOT the default -- unvalidated until it passes + * repeated test_llama_microbench --correctness-only runs (see + * dq8ca_coopmat_variant()'s comment on why a single pass is not proof). + * + * Performs: out[M,N] = dequant(int8_act) * dequant(int_w) (+ bias) + * via coopmat x coopmat -> coopmat on the matrix unit. + * + * Hard preconditions (dbuf4zpg's, plus dbuf4tr's row-major/alignment ones): + * M % WG_TILE_M == 0, N % WG_TILE_N == 0, K % WG_TILE_K == 0, + * group_size % WG_TILE_K == 0, K % 4 == 0, + * WG_TILE_M % MMA_M == 0, WG_TILE_K % MMA_K == 0, + * t_packed_int8_input in kPackedInt8_4W (row-major) layout, + * device exposes coopmatx-> at 16x16x16. + */ + +#version 450 core + +#extension GL_KHR_cooperative_matrix : require +#extension GL_KHR_memory_scope_semantics : require +#extension GL_KHR_shader_subgroup_basic : enable +#extension GL_EXT_shader_explicit_arithmetic_types : require +#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require +// 8-bit SSBO access: A is bound as a scalar int8_t array so that the +// coopMatLoad below has a MATCHING component type (see dbuf4tr's header for +// why the type must match on this driver). +#extension GL_EXT_shader_8bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_EXT_control_flow_attributes : enable + +#define PRECISION ${PRECISION} + +$if WEIGHT_NBITS == 4: + #define WEIGHT_INT4 + +$if HAS_BIAS: + #define HAS_BIAS + +$if WEIGHT_STORAGE == "buffer": + #define WEIGHT_BUFFER + +$if IO_STORAGE == "texture3d": + #define IO_TEXTURE + +layout(std430) buffer; + +#include "common.glslh" + +// Bindings — match add_linear_dqa_qw_node arg order: +// output(0), fp_input(1), packed_int8_input(2), int_input_sums(3 - unused), +// input_scales(4), input_zps(5), packed_weight(6), weight_sums(7), +// weight_scales(8), bias(9). +${layout_declare_tensor(B, "w", "t_output", "half", IO_STORAGE, is_scalar_array=True)} +// t_input is unread here -- the activations arrive already quantized in +// t_packed_int8_input -- but stays declared so the binding layout matches the +// dispatch site. It tracks IO_STORAGE so the two IO tensors stay consistent. +${layout_declare_tensor(B, "r", "t_input", "half", IO_STORAGE, is_scalar_array=False)} +// ROW-MAJOR (kPackedInt8_4W) packed activations, bound as a scalar int8_t +// array (row stride = K int8) -- dbuf4tr's binding, unchanged. The stock +// 4h4w layout dbuf4zpg uses is NOT row-major (component index selects a row, +// non-affine), so it cannot be addressed by any coopMatLoad. +${layout_declare_tensor(B, "r", "t_packed_int8_input", "int8", "buffer", is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_int8_input_sums", "int", "buffer", is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_int8_input_scales", "half", "texture3d")} +${layout_declare_tensor(B, "r", "t_int8_input_zps", "int8", "texture3d")} +${layout_declare_tensor(B, "r", "t_packed_weight", "int", WEIGHT_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_weight_sums", "int", "buffer", is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_weight_scales", "half", "buffer", is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_bias", "half", "buffer", is_scalar_array=True)} + +${layout_declare_ubo(B, "ivec4", "output_sizes")} +${layout_declare_ubo(B, "ivec4", "input_sizes")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +${layout_declare_spec_const(C, "int", "apply_bias", "0")} +// INT4 only; inert (0) for INT8 so the dispatcher's spec list lines up. +${layout_declare_spec_const(C, "int", "K4_per_group", "0")} +${layout_declare_spec_const(C, "int", "num_groups_arg", "0")} +${layout_declare_spec_const(C, "int", "out_N_arg", "0")} + +// Tile geometry +const uint MMA_M = ${MMA_M}; +const uint MMA_N = ${MMA_N}; +const uint MMA_K = ${MMA_K}; + +const uint WG_TILE_M = ${WG_TILE_M}; +const uint WG_TILE_N = ${WG_TILE_N}; +const uint WG_TILE_K = ${WG_TILE_K}; + +const uint SG_GRID_X = ${SG_GRID_X}; +const uint SG_GRID_Y = ${SG_GRID_Y}; +const uint SUBGROUP_SIZE = ${SUBGROUP_SIZE}; +const uint NUM_SUBGROUPS = SG_GRID_X * SG_GRID_Y; +const uint WG_SIZE = NUM_SUBGROUPS * SUBGROUP_SIZE; + +const uint SG_TILE_M = WG_TILE_M / SG_GRID_Y; +const uint SG_TILE_N = WG_TILE_N / SG_GRID_X; +const uint MMAS_PER_SG_M = SG_TILE_M / MMA_M; +const uint MMAS_PER_SG_N = SG_TILE_N / MMA_N; + +const uint A_SLAB_INT8 = WG_TILE_M * MMA_K; +const uint B_USEFUL_U32 = MMA_K / 4u; +// No skew + coalesced write -- dbuf4zpg's B fix, unchanged (this file does +// not touch B staging at all). +const uint B_STRIDE_U32 = B_USEFUL_U32; +const uint B_SLAB_U32 = WG_TILE_N * B_STRIDE_U32; +const uint NUM_K_SLABS = WG_TILE_K / MMA_K; + +const uint A_SLAB_U32 = A_SLAB_INT8 / 4u; +const uint A_STRIDE_U32 = MMA_K / 4u; + +// One ping-pong slice covers all K-slabs of one chunk. +const uint ASH_SLICE_U32 = NUM_K_SLABS * A_SLAB_U32; +const uint BSH_SLICE_U32 = NUM_K_SLABS * B_SLAB_U32; + +// Double-buffered MMA operand staging. +shared uint Ash_int8[2u * ASH_SLICE_U32]; +shared uint Bsh_int8[2u * BSH_SLICE_U32]; + +// Per-WG-tile-row activation params (loaded ONCE at WG start; constant +// across groups). +shared int izp_sh[WG_TILE_M]; // int32 (cast from int8 source) for broadcast +shared float ifs_sh[WG_TILE_M]; // float32 (cast from fp16 source) for broadcast + +// Per-(group, output-channel) weight params, ping-ponged by group parity. +// (For per-channel INT8 only slice 0 is ever used.) +shared float wsc_sh[2u * WG_TILE_N]; +// SUM_g wsc[g][n]*wsum[g][n] per output channel -- weight-side only, so it is +// accumulated once in the prologue. dbuf4zpg's zp-hoist, unchanged. +shared float wcorr_sh[WG_TILE_N]; + +#ifdef HAS_BIAS +shared float bias_sh[WG_TILE_N]; +#endif + +#ifdef IO_TEXTURE +// Result staging for the imageStore epilogue, mirroring the fp16 kernel: +// SG_GRID_Y bands of MMA_M rows, each WG_TILE_N wide, row-major. A full +// WG_TILE_M x WG_TILE_N buffer would cost SG_GRID_Y/MMAS_PER_SG_M x more LDS +// and wreck occupancy. float16_t-typed because coopMatStore needs it. +const uint CSH_ROWS = SG_GRID_Y * MMA_M; +shared float16_t Csh_out[CSH_ROWS * WG_TILE_N]; +#endif + +// Running fp32 accumulator (across all groups). +coopmat + result[MMAS_PER_SG_M][MMAS_PER_SG_N]; + +// Per-group int32 MMA accumulator. +coopmat + accum_int32[MMAS_PER_SG_M][MMAS_PER_SG_N]; + + +// Byte-parallel int4 -> int8 widening. dbuf4zpg's, unchanged (B-side only). +// +// The four nibbles this shader needs from one packed uint are ALREADY one per +// byte (bits 3:0 of each byte for parity 0, bits 7:4 for parity 1), so all four +// can be widened at once instead of with a per-nibble +// shift/mask/bias-subtract/mask chain. +// +// For v in [0,15] the biased value is v-8. `v ^ 8` is exactly the 4-bit +// two's-complement pattern of v-8, because -8 == +8 (mod 16): +// v=0 -> 0x8 -> -8 v=7 -> 0xF -> -1 +// v=8 -> 0x0 -> 0 v=15 -> 0x7 -> +7 +// so the only remaining work is sign-extending bit 3 into bits 7:4 per byte. +// `sgn * 0x1E` does that with no cross-byte carry: 0x08 * 0x1E == 0xF0 exactly, +// and sgn is at most 0x08080808 so the product is at most 0xF0F0F0F0. +// +// A naive `nib - 0x08080808` would NOT work -- it borrows across byte lanes +// whenever a nibble is < 8. Shifts must be on uint, not int, so the >> is +// logical rather than arithmetic. +// +// ~5 ops per 4 weights vs ~22 for the per-nibble chain; bit-identical output. +uint widen_nibbles(const uint w, const uint parity) { + const uint nib = (parity == 0u) ? (w & 0x0F0F0F0Fu) : ((w >> 4u) & 0x0F0F0F0Fu); + const uint p = nib ^ 0x08080808u; + const uint sgn = p & 0x08080808u; + return p | (sgn * 0x1Eu); +} + +void main() { + const uvec2 tileID = uvec2(gl_WorkGroupID.xy); + const uvec2 warpInTile = uvec2( + gl_SubgroupID % SG_GRID_X, + gl_SubgroupID / SG_GRID_X); + + const uint K = uint(input_sizes.x); + const uint N = uint(output_sizes.x); + const uint N4 = (N + 3u) / 4u; + const uint nblocks_x_A = (K + 3u) >> 2u; + // A row stride in INT8 elements (dbuf4tr's binding is row-major int8, not + // the 4h4w ivec4 block layout dbuf4zpg used -- so A addressing below is in + // int8 elements, not int, and derived from nblocks_x_A so it matches the + // packer's `m_row * K4 + k4` addressing exactly (K % 4 == 0 makes them + // equal to K directly). + const uint a_row_stride_i8 = nblocks_x_A * 4u; + +#ifdef WEIGHT_INT4 + const uint num_groups = uint(num_groups_arg); + const uint CHUNKS_PER_GROUP = uint(K4_per_group) * 4u / WG_TILE_K; +#else + // Per-channel: a single quant "group" spanning all of K. The nested + // groups x chunks loop below collapses to a flat chunk loop, the wsum/wsc + // ping-pong never crosses a boundary, and the epilog runs exactly once. + const uint num_groups = 1u; + const uint CHUNKS_PER_GROUP = uint(num_groups_arg); +#endif + const uint num_chunks = num_groups * CHUNKS_PER_GROUP; + + const uint tile_m_start = WG_TILE_M * tileID.y; + const uint tile_n_start = WG_TILE_N * tileID.x; + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + result[i][j] = coopmat(0.0); + accum_int32[i][j] = coopmat(0); + } + } + + const uint K_BLOCKS_PER_CHUNK = WG_TILE_K >> 2u; + + // --- A staging tile map (dbuf4tr's, unmodified): one MMA_M x MMA_K coopmat + // tile per subgroup per slot, dealt round-robin across the + // NUM_SUBGROUPS subgroups so every subgroup participates. Replaces + // dbuf4zpg's per-thread (m4, k4) map / a_active guard entirely -- see + // design.md D3 for why this is reused as-is, not re-derived. + const uint A_TILES_M = WG_TILE_M / MMA_M; + const uint A_TILES_K = WG_TILE_K / MMA_K; // == NUM_K_SLABS + const uint NUM_A_TILES = A_TILES_M * A_TILES_K; + const uint A_TILES_PER_SG = (NUM_A_TILES + NUM_SUBGROUPS - 1u) / NUM_SUBGROUPS; + +#ifdef WEIGHT_INT4 + // --- B staging thread map: (block, col) slots; each slot extracts one + // ColumnMajor LDS uint (4 K-contiguous sign-extended int8) --- + const uint B_TOTAL_SLOTS = K_BLOCKS_PER_CHUNK * WG_TILE_N; + const uint B_SLOTS_PER_THREAD = B_TOTAL_SLOTS / WG_SIZE; + const uint N8_PER_TILE = WG_TILE_N >> 3u; +#else + // --- B staging thread map: one (k4, n4) ivec4 block per active thread --- + const uint B_FETCH_SLOTS = K_BLOCKS_PER_CHUNK * (WG_TILE_N >> 2u); + const uint N4_PER_TILE = WG_TILE_N >> 2u; + const uint b_k4_in_chunk = gl_LocalInvocationID.x / N4_PER_TILE; + const uint b_n_uint_col = gl_LocalInvocationID.x % N4_PER_TILE; + const bool b_active = gl_LocalInvocationID.x < B_FETCH_SLOTS; +#endif + + // ===== INTERVENTION F: hoist loop-invariant B staging index math ===== + // dbuf4zpg's, unchanged -- B staging is untouched by this file's A-staging + // swap. See dbuf4zpg's header for the full rationale (ablation-attributed + // -16.8% block, mostly index arithmetic not memory traffic). +#ifdef WEIGHT_INT4 + // ===== bw2 DIAGNOSTIC: 2 consecutive uints per thread, ALL 256 threads + // active (no idling). Discriminates "4-wide ownership + idle threads" from + // "consecutive-write ownership per se" as the cause of the num_groups==2 + // failure seen in bw/bwr. + const uint B_PAIR = 2u; + const uint b_lin0 = gl_LocalInvocationID.x * B_PAIR; // first uint index + const uint b_slab = b_lin0 / B_SLAB_U32; + const uint b_inslab = b_lin0 % B_SLAB_U32; + const uint b_col = b_inslab / B_STRIDE_U32; + const uint b_k4lo = b_inslab % B_STRIDE_U32; + const uint b_rem = b_col & 7u; + const uint b_comp1 = b_rem & 3u; + const uint b_par1 = b_rem >> 2u; + const uint b_n8blk1 = (tile_n_start >> 3u) + (b_col >> 3u); + const uint b_k4base = b_slab * (MMA_K >> 2u) + b_k4lo; + const uint b_lds_base = b_lin0; +#endif + + // Prefetch temp registers. temp_A is a coopmat array (dbuf4tr's A-staging + // technique); indices into it are [[unroll]]-resolved compile-time + // constants, never dynamic -- dynamic indexing of a coopmat array is + // exactly the construct the Xclipse/AMD-PAL compiler has miscompiled + // before. + coopmat + temp_A[A_TILES_PER_SG]; +#ifdef WEIGHT_INT4 + ivec4 temp_B[2u]; // bw2: one per k4 of the owned pair + float temp_wsc; +#else + ivec4 temp_B; +#endif + + // ========================================================= + // PROLOGUE + // ========================================================= + if (gl_LocalInvocationID.x < (WG_TILE_M >> 2u)) { + const uint m4 = (tile_m_start >> 2u) + gl_LocalInvocationID.x; + const vec4 sc = vec4(texelFetch(t_int8_input_scales, ivec3(m4, 0, 0), 0)); + const ivec4 zp = texelFetch(t_int8_input_zps, ivec3(m4, 0, 0), 0); + const uint base = gl_LocalInvocationID.x * 4u; + ifs_sh[base + 0u] = sc.x; ifs_sh[base + 1u] = sc.y; + ifs_sh[base + 2u] = sc.z; ifs_sh[base + 3u] = sc.w; + izp_sh[base + 0u] = zp.x; izp_sh[base + 1u] = zp.y; + izp_sh[base + 2u] = zp.z; izp_sh[base + 3u] = zp.w; + } + // Group 0 weight scales -> slice 0, and the hoisted weight-side correction + // SUM_g wsc[g][n]*wsum[g][n] accumulated across ALL groups. dbuf4zpg's + // zp-hoist, unchanged. + if (gl_LocalInvocationID.x < WG_TILE_N) { + const uint n_idx = tile_n_start + gl_LocalInvocationID.x; + f16vec4 sv0 = t_weight_scales[n_idx >> 2u]; + wsc_sh[gl_LocalInvocationID.x] = float(sv0[n_idx & 3u]); + + float corr = 0.0; + for (uint g = 0; g < num_groups; ++g) { + f16vec4 sv = t_weight_scales[g * N4 + (n_idx >> 2u)]; + corr += float(sv[n_idx & 3u]) * float(t_weight_sums[g * N + n_idx]); + } + wcorr_sh[gl_LocalInvocationID.x] = corr; + } + memoryBarrierShared(); + barrier(); + + // NOTE: dbuf4zpg builds izp_bcast/ifs_bcast AFTER the group loop, not here + // -- that is the register-pressure saving zp-hoist buys. Unchanged. + + // dbuf4: prefetch chunk 0 into temp registers, THEN store to slice 0 (no + // barrier here -- the main loop's first iteration barriers before + // reading slice 0). + // + // A staging (dbuf4tr's technique): per-subgroup coopMatLoad straight from + // the row-major global buffer. + [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { + const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; + if (t < NUM_A_TILES) { + const uint tm = t / A_TILES_K; + const uint tk = t % A_TILES_K; + coopMatLoad( + temp_A[s], t_packed_int8_input, + (tile_m_start + tm * MMA_M) * a_row_stride_i8 + tk * MMA_K, + a_row_stride_i8, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#ifdef WEIGHT_INT4 + [[unroll]] for (uint kk = 0; kk < 2u; ++kk) { +#ifdef WEIGHT_BUFFER + temp_B[kk] = t_packed_weight[(b_n8blk1 * nblocks_x_A) + b_k4base + kk]; +#else + temp_B[kk] = texelFetch(t_packed_weight, ivec2(b_k4base + kk, b_n8blk1), 0); +#endif + } +#else + if (b_active) { + const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; +#ifdef WEIGHT_BUFFER + temp_B = t_packed_weight[(b_k4_in_chunk * N4) + block_x_w]; +#else + temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, b_k4_in_chunk), 0); +#endif + } +#endif + { + // store chunk 0 -> slice 0 + // A staging (dbuf4tr's technique): coopMatStore into the same Ash_int8 + // slot layout dbuf4zpg's scalar scatter used to write. + [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { + const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; + if (t < NUM_A_TILES) { + const uint tm = t / A_TILES_K; + const uint tk = t % A_TILES_K; + coopMatStore( + temp_A[s], Ash_int8, + tk * A_SLAB_U32 + (tm * MMA_M) * A_STRIDE_U32, + A_STRIDE_U32, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#ifdef WEIGHT_INT4 + [[unroll]] for (uint kk = 0; kk < 2u; ++kk) { + Bsh_int8[b_lds_base + kk] = + widen_nibbles(uint(temp_B[kk][b_comp1]), b_par1); + } +#else + if (b_active) { + const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); + const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); + const uint n_col_base = b_n_uint_col * 4u; + [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { + Bsh_int8[slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab] = + uint(temp_B[n_in_blk]); + } + } +#endif + } + + // ========================================================= + // MAIN LOOP (dbuf4) — nested groups x chunks (kept nested; flattening it + // with a conditional coopmat epilog crashes the Xclipse PAL compiler at + // large spec-resolved trip counts). One barrier per chunk. Chunk + // iteration (global index `chunk`): + // 1. barrier — A/B slice (chunk%2) fully written; on the first chunk + // of group g, wsc slice (g%2) is too. + // 2. prefetch — chunk+1 (A tiles, B blocks) into temp; when chunk+1 + // starts a new group, also its wsc element. Skipped + // entirely on the final chunk. + // 3. int8 MMA — on slice (chunk%2) into accum_int32. + // 4. store — temp -> A/B slice ((chunk+1)%2), unpacking the weight; + // on a group boundary, wsc -> slice ((g+1)%2). + // The group epilog runs unconditionally at the tail of each group. + // ========================================================= + uint chunk = 0; + for (uint group_i = 0; group_i < num_groups; ++group_i) { + for (uint inner = 0; inner < CHUNKS_PER_GROUP; ++inner, ++chunk) { + const bool has_next = chunk + 1u < num_chunks; + const bool group_crossing = has_next && (inner + 1u == CHUNKS_PER_GROUP); + const uint cur_a = (chunk % 2u) * ASH_SLICE_U32; + const uint cur_b = (chunk % 2u) * BSH_SLICE_U32; + const uint nxt_a = ((chunk + 1u) % 2u) * ASH_SLICE_U32; + const uint nxt_b = ((chunk + 1u) % 2u) * BSH_SLICE_U32; + + // coopmat-lds-fence: barrier() alone does NOT order shared stores against a + // subsequent coopMatLoad on the M51 Xclipse/AMD-PAL driver -- symptom is one + // stale MMA_M-row band of A, all columns, ~2.5% of runs, no crash (observed + // 2026-09-02 in sdpa_compute_out_coopmat.glsl). Measured cost of the fence: + // none (see this change's results). See memory + // `coopmat-lds-needs-explicit-memorybarriershared`. + memoryBarrierShared(); + barrier(); + + // --- 2. prefetch chunk+1 -> temp --- + if (has_next) { + const uint chunkK_nxt = (chunk + 1u) * WG_TILE_K; + // A staging (dbuf4tr's technique): coopMatLoad straight from global. + [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { + const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; + if (t < NUM_A_TILES) { + const uint tm = t / A_TILES_K; + const uint tk = t % A_TILES_K; + coopMatLoad( + temp_A[s], t_packed_int8_input, + (tile_m_start + tm * MMA_M) * a_row_stride_i8 + chunkK_nxt + + tk * MMA_K, + a_row_stride_i8, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#ifdef WEIGHT_INT4 + [[unroll]] for (uint kk = 0; kk < 2u; ++kk) { + const uint k4_blk = (chunkK_nxt >> 2u) + b_k4base + kk; +#ifdef WEIGHT_BUFFER + temp_B[kk] = t_packed_weight[(b_n8blk1 * nblocks_x_A) + k4_blk]; +#else + temp_B[kk] = texelFetch(t_packed_weight, ivec2(k4_blk, b_n8blk1), 0); +#endif + } + if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { + const uint n_idx = tile_n_start + gl_LocalInvocationID.x; + f16vec4 sv = t_weight_scales[(group_i + 1u) * N4 + (n_idx >> 2u)]; + temp_wsc = float(sv[n_idx & 3u]); + } +#else + if (b_active) { + const uint block_y_w = (chunkK_nxt >> 2u) + b_k4_in_chunk; + const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; +#ifdef WEIGHT_BUFFER + temp_B = t_packed_weight[(block_y_w * N4) + block_x_w]; +#else + temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, block_y_w), 0); +#endif + } +#endif + } + + // --- 3. int8 MMA on the cur slice --- + [[unroll]] for (uint k = 0; k < NUM_K_SLABS; ++k) { + const uint slab_a_base_u32 = cur_a + k * A_SLAB_U32; + const uint slab_b_base_u32 = cur_b + k * B_SLAB_U32; + + coopmat matA[MMAS_PER_SG_M]; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + const uint row_a = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + coopMatLoad( + matA[i], Ash_int8, + slab_a_base_u32 + row_a * A_STRIDE_U32, + A_STRIDE_U32, + gl_CooperativeMatrixLayoutRowMajor); + } + + coopmat matB; + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint col_b = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopMatLoad( + matB, Bsh_int8, + slab_b_base_u32 + col_b * B_STRIDE_U32, + B_STRIDE_U32, + gl_CooperativeMatrixLayoutColumnMajor); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + accum_int32[i][j] = coopMatMulAdd(matA[i], matB, accum_int32[i][j]); + } + } + } + + // --- 4. store temp (chunk+1) -> nxt slice --- + if (has_next) { + // A staging (dbuf4tr's technique): coopMatStore into the nxt slice. + [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { + const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; + if (t < NUM_A_TILES) { + const uint tm = t / A_TILES_K; + const uint tk = t % A_TILES_K; + coopMatStore( + temp_A[s], Ash_int8, + nxt_a + tk * A_SLAB_U32 + (tm * MMA_M) * A_STRIDE_U32, + A_STRIDE_U32, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#ifdef WEIGHT_INT4 + [[unroll]] for (uint kk = 0; kk < 2u; ++kk) { + Bsh_int8[nxt_b + b_lds_base + kk] = + widen_nibbles(uint(temp_B[kk][b_comp1]), b_par1); + } + if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { + const uint wbase_nxt = ((group_i + 1u) % 2u) * WG_TILE_N; + wsc_sh[wbase_nxt + gl_LocalInvocationID.x] = temp_wsc; + } +#else + if (b_active) { + const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); + const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); + const uint n_col_base = b_n_uint_col * 4u; + [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { + Bsh_int8[nxt_b + slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab] = + uint(temp_B[n_in_blk]); + } + } +#endif + } + } // chunks + + // --- Group epilog: scale-only accumulate, reset accum --- + // dbuf4zpg's, unchanged. Just result += float(acc) * wsc. The + // zero-point subtract and the ifs multiply are hoisted out of the group + // loop (applied once below). + { + const uint wbase = (group_i % 2u) * WG_TILE_N; + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + + coopmat wsc_bcast; + coopMatLoad( + wsc_bcast, wsc_sh, + wbase + local_n_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutRowMajor); + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + result[i][j] += + coopmat( + accum_int32[i][j]) * wsc_bcast; + accum_int32[i][j] = coopmat(0); + } + } + } + } // groups + + // --- Hoisted correction, applied ONCE: -------------------------------- + // result = ifs * ( result - izp * SUM_g wsc_g*wsum_g ) + // dbuf4zpg's, unchanged. izp/ifs are loaded here rather than before the + // group loop so they are not live across it. + { + coopmat + izpf_bcast[MMAS_PER_SG_M]; + coopmat + ifs_bcast[MMAS_PER_SG_M]; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + const uint local_m_base = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + coopmat izp_i; + coopMatLoad( + izp_i, izp_sh, local_m_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutColumnMajor); + izpf_bcast[i] = + coopmat(izp_i); + coopMatLoad( + ifs_bcast[i], ifs_sh, local_m_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutColumnMajor); + } + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopmat wcorr_bcast; + coopMatLoad( + wcorr_bcast, wcorr_sh, local_n_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutRowMajor); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + result[i][j] = ifs_bcast[i] * (result[i][j] - izpf_bcast[i] * wcorr_bcast); + } + } + } + + // --- Bias (optional) --- +#ifdef HAS_BIAS + if (apply_bias > 0) { + for (uint t = gl_LocalInvocationID.x; t < WG_TILE_N; t += WG_SIZE) { + bias_sh[t] = float(t_bias[tile_n_start + t]); + } + memoryBarrierShared(); + barrier(); + } +#endif + + // --- Store result tile --- + // N for the store address math MUST come from the spec constant, not the + // sizes UBO (see out_N_arg above). +#ifdef IO_TEXTURE + // Epilogue iteration i drains accumulator row-block i from EVERY subgroup + // into Csh_out at once, so the SG_GRID_Y bands it holds are disjoint global + // row ranges; the whole workgroup then imageStores them. lr / MMA_M is the + // writing subgroup's warpInTile.y, so the global row reproduces the buffer + // path's gi exactly. + // + // PORTABILITY NOTE: the barrier() in the loop body keeps this loop rolled + // despite [[unroll]], so result[i][j] IS dynamically indexed. Coopmat arrays + // are opaque per-lane storage and dynamic indexing is exactly the construct + // the Xclipse/AMD-PAL compiler has broken before -- check this first if the + // texture variants miscompile on M51. + const uint CSH_TEXELS_PER_ROW = WG_TILE_N / 4u; + const uint CSH_TEXELS = CSH_ROWS * CSH_TEXELS_PER_ROW; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + // Guards Csh_out against the previous iteration's readers. Inert on i == 0 + // but must stay unconditional to remain workgroup-uniform. + // coopmat-lds-fence: barrier() alone does NOT order shared stores against a + // subsequent coopMatLoad on the M51 Xclipse/AMD-PAL driver -- symptom is one + // stale MMA_M-row band of A, all columns, ~2.5% of runs, no crash (observed + // 2026-09-02 in sdpa_compute_out_coopmat.glsl). Measured cost of the fence: + // none (see this change's results). See memory + // `coopmat-lds-needs-explicit-memorybarriershared`. + memoryBarrierShared(); + barrier(); + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { +#ifdef HAS_BIAS + if (apply_bias > 0) { + const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopmat bias_tile; + coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); + result[i][j] += bias_tile; + } +#endif + coopmat out_tile = + coopmat(result[i][j]); + coopMatStore( + out_tile, Csh_out, + warpInTile.y * MMA_M * WG_TILE_N + + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j), + WG_TILE_N, + gl_CooperativeMatrixLayoutRowMajor); + } + memoryBarrierShared(); + barrier(); + + for (uint t = gl_LocalInvocationID.x; t < CSH_TEXELS; t += WG_SIZE) { + const uint lr = t / CSH_TEXELS_PER_ROW; + const uint lc4 = t % CSH_TEXELS_PER_ROW; + const uint m = + tile_m_start + (lr / MMA_M) * SG_TILE_M + i * MMA_M + (lr % MMA_M); + const uint base = lr * WG_TILE_N + lc4 * 4u; + imageStore( + t_output, + ivec3(tile_n_start / 4u + lc4, m, 0), + vec4( + float(Csh_out[base]), + float(Csh_out[base + 1u]), + float(Csh_out[base + 2u]), + float(Csh_out[base + 3u]))); + } + } +#else + const uint N_out = uint(out_N_arg); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint gi = tile_m_start + MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + const uint gj = tile_n_start + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + +#ifdef HAS_BIAS + if (apply_bias > 0) { + const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopmat bias_tile; + coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); + result[i][j] += bias_tile; + } +#endif + + coopmat out_tile = + coopmat(result[i][j]); + coopMatStore( + out_tile, t_output, + gi * N_out + gj, N_out, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#endif // IO_TEXTURE +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw2.yaml b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw2.yaml new file mode 100644 index 00000000000..c51f9cb0359 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw2.yaml @@ -0,0 +1,60 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# "zpg" + "-tr" combination: tsweep_dbuf4zpg with its per-thread scalar +# A-staging replaced by tsweep_dbuf4tr's coopMat-mediated A staging. B +# staging, zp-hoist, byte-parallel nibble widening are all dbuf4zpg's, +# unchanged -- only A-staging differs. Requires t_packed_int8_input in the +# ROW-MAJOR kPackedInt8_4W layout (same requirement as tsweep_dbuf4tr). +# +# PROMOTED 2026-09-01 as the shipped default -- see QuantizedLinear.cpp's +# dq8ca_coopmat_variant() for the full validation record (10/10 buffer, 6/6 +# texture3d correctness across 1B/3B/8B; +4.2% real e2e prefill; 46.50% -> +# 49.94% efficiency on 8B). Also selectable explicitly via +# ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4zpgbw2_txkgs. +# See openspec/changes/archive/2026-08-31-dq8ca-tr-staged-a-on-zpg. +# +# Single seed tile: the prior dbuf4zpg default's own tile (t128x64k32g42s32). +# A re-sweep against this shader's own (lower) register-pressure profile was +# run as a follow-up (coopmat-tr-tilesweep-4w-port) and found no better tile +# -- this remains the best known geometry. + +linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw2: + parameter_names_with_default_values: + PRECISION: highp + HAS_BIAS: false + IO_STORAGE: buffer + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + MMA_M: 16 + MMA_N: 16 + MMA_K: 16 + WG_TILE_M: 64 + WG_TILE_N: 32 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + shader_variants: + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw2_t128x64k32g42s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw2_t128x64k32g42s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw3.glsl b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw3.glsl new file mode 100644 index 00000000000..5f71b6ed5e1 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw3.glsl @@ -0,0 +1,798 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +/* + * "zpg" + "-tr" combination: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpg.glsl + * with its per-thread scalar A-staging replaced by + * linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr.glsl's coopMat-mediated A + * staging, PLUS "bw3": shared B retyped to uvec2 with dbuf4zpgbw2's + * all-threads-active ownership, so each thread issues ONE ds_write_b64 for its + * 2-uint pair. A staging, the quantization math and every safeguard are + * dbuf4zpgtr's, unchanged. This is + * an ADDITIVE combination, not a redesign: every non-A-staging block below + * (B staging: coalesced write, no skew; zp-hoist: izp/ifs applied once after + * the group loop via wcorr_sh; byte-parallel nibble widening; static + * A_ALWAYS_ACTIVE branch elision -- N/A here, see below; group epilog; + * bias/store epilogue) is byte-identical to dbuf4zpg's. Only the A-staging + * block (prologue load+store, main-loop prefetch+store) is dbuf4tr's, + * verbatim. + * + * dbuf4zpg's per-thread A staging used an `a_active` guard (statically always + * true when A_ACTIVE_THREADS == WG_SIZE, via the A_MAP_FULL-gated + * A_ALWAYS_ACTIVE macro). dbuf4tr's per-SUBGROUP tile map has no equivalent + * concept -- every subgroup participates via a `t < NUM_A_TILES` guard that + * depends only on gl_SubgroupID, not gl_LocalInvocationID.x -- so + * A_MAP_FULL/A_ALWAYS_ACTIVE is dropped entirely in this file; it would be + * dead code for the new A-staging block. + * + * Rationale for combining this way (not the reverse) and why this is worth + * building at all: see this change's design.md D0-D3. In short -- the only + * existing measurement of dbuf4tr's A-staging technique (28.72-30.51%, + * dq8ca-arch-redesign) was taken against dbuf4tr's own pre-zpg baseline + * (old B skew, no byte-parallel widening, no branch elision) -- a materially + * weaker shader than the 46.49-46.50% dbuf4zpg this file now combines it + * with. This file exists to answer whether that combination performs + * differently now that register pressure is already reduced. + * + * A staging (the actual delta from dbuf4zpg): + * dbuf4zpg: per-thread (m4, k4) ivec4 fetch, hoisted a_lds_off0/a_glb_row; + * only A_ACTIVE_THREADS invocations participate, each scattering + * 4 rows into Ash_int8 with 4 scalar stores. + * this file: per-SUBGROUP MMA_M x MMA_K tile fetch via coopMatLoad straight + * from a ROW-MAJOR (kPackedInt8_4W) int8 activation buffer, then + * coopMatStore into the same Ash_int8 slot -- dbuf4tr's mapping, + * unmodified (not re-derived; see design.md D3). + * + * t_packed_int8_input is therefore bound the same way dbuf4tr binds it: a + * SCALAR int8_t array in the kPackedInt8_4W layout (plain row-major int8, + * row stride K), produced by quantize_and_pack_4w_with_group_sums.glsl. + * QuantizedLinear.cpp's dq8ca_variant_wants_rowmajor_a() must recognize this + * file's variant token (tsweep_dbuf4zpgtr_t...) the same way it already + * recognizes tsweep_dbuf4tr_t/trm_t/trd_t, so graph-build time (packer + * selection) and dispatch time (kernel selection) cannot disagree. + * + * B CANNOT be coopmat-staged (int4 nibble unpack; a coopmat's per-lane layout + * is opaque to hand-assembly from unpacked registers) -- unchanged from both + * parent files. B staging below is dbuf4zpg's byte-parallel, coalesced, + * no-skew version, untouched. + * + * The loop structure is dbuf4's (both parents share it), unchanged: + * prologue: prefetch chunk 0 -> temp, store to slice 0 (no barrier) + * per iter: barrier -> prefetch(next) -> MMA(cur) -> store(next) + * kept nested (groups x chunks) with an unconditional group epilog -- + * flattening it crashes the Xclipse PAL compiler at large spec-resolved trip + * counts (see dbuf2's own header). + * + * Selected via + * ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4zpgtr_txkgs<32|64> + * (QuantizedLinear.cpp), additive to the tsweep_dbuf4zpg_t..., tsweep_dbuf4tr_t... + * and tsweep_t... namespaces. NOT the default -- unvalidated until it passes + * repeated test_llama_microbench --correctness-only runs (see + * dq8ca_coopmat_variant()'s comment on why a single pass is not proof). + * + * Performs: out[M,N] = dequant(int8_act) * dequant(int_w) (+ bias) + * via coopmat x coopmat -> coopmat on the matrix unit. + * + * Hard preconditions (dbuf4zpg's, plus dbuf4tr's row-major/alignment ones): + * M % WG_TILE_M == 0, N % WG_TILE_N == 0, K % WG_TILE_K == 0, + * group_size % WG_TILE_K == 0, K % 4 == 0, + * WG_TILE_M % MMA_M == 0, WG_TILE_K % MMA_K == 0, + * t_packed_int8_input in kPackedInt8_4W (row-major) layout, + * device exposes coopmatx-> at 16x16x16. + */ + +#version 450 core + +#extension GL_KHR_cooperative_matrix : require +#extension GL_KHR_memory_scope_semantics : require +#extension GL_KHR_shader_subgroup_basic : enable +#extension GL_EXT_shader_explicit_arithmetic_types : require +#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require +// 8-bit SSBO access: A is bound as a scalar int8_t array so that the +// coopMatLoad below has a MATCHING component type (see dbuf4tr's header for +// why the type must match on this driver). +#extension GL_EXT_shader_8bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_EXT_control_flow_attributes : enable + +#define PRECISION ${PRECISION} + +$if WEIGHT_NBITS == 4: + #define WEIGHT_INT4 + +$if HAS_BIAS: + #define HAS_BIAS + +$if WEIGHT_STORAGE == "buffer": + #define WEIGHT_BUFFER + +$if IO_STORAGE == "texture3d": + #define IO_TEXTURE + +layout(std430) buffer; + +#include "common.glslh" + +// Bindings — match add_linear_dqa_qw_node arg order: +// output(0), fp_input(1), packed_int8_input(2), int_input_sums(3 - unused), +// input_scales(4), input_zps(5), packed_weight(6), weight_sums(7), +// weight_scales(8), bias(9). +${layout_declare_tensor(B, "w", "t_output", "half", IO_STORAGE, is_scalar_array=True)} +// t_input is unread here -- the activations arrive already quantized in +// t_packed_int8_input -- but stays declared so the binding layout matches the +// dispatch site. It tracks IO_STORAGE so the two IO tensors stay consistent. +${layout_declare_tensor(B, "r", "t_input", "half", IO_STORAGE, is_scalar_array=False)} +// ROW-MAJOR (kPackedInt8_4W) packed activations, bound as a scalar int8_t +// array (row stride = K int8) -- dbuf4tr's binding, unchanged. The stock +// 4h4w layout dbuf4zpg uses is NOT row-major (component index selects a row, +// non-affine), so it cannot be addressed by any coopMatLoad. +${layout_declare_tensor(B, "r", "t_packed_int8_input", "int8", "buffer", is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_int8_input_sums", "int", "buffer", is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_int8_input_scales", "half", "texture3d")} +${layout_declare_tensor(B, "r", "t_int8_input_zps", "int8", "texture3d")} +${layout_declare_tensor(B, "r", "t_packed_weight", "int", WEIGHT_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_weight_sums", "int", "buffer", is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_weight_scales", "half", "buffer", is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_bias", "half", "buffer", is_scalar_array=True)} + +${layout_declare_ubo(B, "ivec4", "output_sizes")} +${layout_declare_ubo(B, "ivec4", "input_sizes")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +${layout_declare_spec_const(C, "int", "apply_bias", "0")} +// INT4 only; inert (0) for INT8 so the dispatcher's spec list lines up. +${layout_declare_spec_const(C, "int", "K4_per_group", "0")} +${layout_declare_spec_const(C, "int", "num_groups_arg", "0")} +${layout_declare_spec_const(C, "int", "out_N_arg", "0")} + +// Tile geometry +const uint MMA_M = ${MMA_M}; +const uint MMA_N = ${MMA_N}; +const uint MMA_K = ${MMA_K}; + +const uint WG_TILE_M = ${WG_TILE_M}; +const uint WG_TILE_N = ${WG_TILE_N}; +const uint WG_TILE_K = ${WG_TILE_K}; + +const uint SG_GRID_X = ${SG_GRID_X}; +const uint SG_GRID_Y = ${SG_GRID_Y}; +const uint SUBGROUP_SIZE = ${SUBGROUP_SIZE}; +const uint NUM_SUBGROUPS = SG_GRID_X * SG_GRID_Y; +const uint WG_SIZE = NUM_SUBGROUPS * SUBGROUP_SIZE; + +const uint SG_TILE_M = WG_TILE_M / SG_GRID_Y; +const uint SG_TILE_N = WG_TILE_N / SG_GRID_X; +const uint MMAS_PER_SG_M = SG_TILE_M / MMA_M; +const uint MMAS_PER_SG_N = SG_TILE_N / MMA_N; + +const uint A_SLAB_INT8 = WG_TILE_M * MMA_K; +const uint B_USEFUL_U32 = MMA_K / 4u; +// No skew + coalesced write -- dbuf4zpg's B fix, unchanged (this file does +// not touch B staging at all). +const uint B_STRIDE_U32 = B_USEFUL_U32; +const uint B_SLAB_U32 = WG_TILE_N * B_STRIDE_U32; +const uint NUM_K_SLABS = WG_TILE_K / MMA_K; + +const uint A_SLAB_U32 = A_SLAB_INT8 / 4u; +const uint A_STRIDE_U32 = MMA_K / 4u; + +// One ping-pong slice covers all K-slabs of one chunk. +const uint ASH_SLICE_U32 = NUM_K_SLABS * A_SLAB_U32; +const uint BSH_SLICE_U32 = NUM_K_SLABS * B_SLAB_U32; + +// Double-buffered MMA operand staging. +shared uint Ash_int8[2u * ASH_SLICE_U32]; +// ===== bw3: shared B retyped to uvec2 ===== +// The ONLY combination that gets both a genuinely wide LDS store and the +// safe ownership. A uvec4 element would be a whole column-slab, forcing +// 4-wide-per-thread ownership with half the workgroup idle -- which is exactly +// the pattern that fails the num_groups==2 shapes (see dbuf4zpgbw/bwr/bv4). +// A uvec2 element is 2 uints, so BSH_SLICE_V2 == WG_SIZE at the shipped +// geometry: one element per thread, ALL 256 threads active (dbuf4zpgbw2's +// proven-correct ownership), one ds_write_b64 each. +// Packing ratio for the compute-side coopMatLoad is 8x (int8 coopmat from +// uvec2[]); dbuf4zpgbw2 showed adjacent scalar writes do NOT get merged by the +// compiler, so retyping is the only way to actually emit a wide store. +const uint B_STRIDE_V2 = B_STRIDE_U32 / 2u; // uvec2 per column +const uint B_SLAB_V2 = B_SLAB_U32 / 2u; +const uint BSH_SLICE_V2 = BSH_SLICE_U32 / 2u; +const uint V2_PER_THREAD = BSH_SLICE_V2 / WG_SIZE; +shared uvec2 Bsh_v2[2u * BSH_SLICE_V2]; + +// Per-WG-tile-row activation params (loaded ONCE at WG start; constant +// across groups). +shared int izp_sh[WG_TILE_M]; // int32 (cast from int8 source) for broadcast +shared float ifs_sh[WG_TILE_M]; // float32 (cast from fp16 source) for broadcast + +// Per-(group, output-channel) weight params, ping-ponged by group parity. +// (For per-channel INT8 only slice 0 is ever used.) +shared float wsc_sh[2u * WG_TILE_N]; +// SUM_g wsc[g][n]*wsum[g][n] per output channel -- weight-side only, so it is +// accumulated once in the prologue. dbuf4zpg's zp-hoist, unchanged. +shared float wcorr_sh[WG_TILE_N]; + +#ifdef HAS_BIAS +shared float bias_sh[WG_TILE_N]; +#endif + +#ifdef IO_TEXTURE +// Result staging for the imageStore epilogue, mirroring the fp16 kernel: +// SG_GRID_Y bands of MMA_M rows, each WG_TILE_N wide, row-major. A full +// WG_TILE_M x WG_TILE_N buffer would cost SG_GRID_Y/MMAS_PER_SG_M x more LDS +// and wreck occupancy. float16_t-typed because coopMatStore needs it. +const uint CSH_ROWS = SG_GRID_Y * MMA_M; +shared float16_t Csh_out[CSH_ROWS * WG_TILE_N]; +#endif + +// Running fp32 accumulator (across all groups). +coopmat + result[MMAS_PER_SG_M][MMAS_PER_SG_N]; + +// Per-group int32 MMA accumulator. +coopmat + accum_int32[MMAS_PER_SG_M][MMAS_PER_SG_N]; + + +// Byte-parallel int4 -> int8 widening. dbuf4zpg's, unchanged (B-side only). +// +// The four nibbles this shader needs from one packed uint are ALREADY one per +// byte (bits 3:0 of each byte for parity 0, bits 7:4 for parity 1), so all four +// can be widened at once instead of with a per-nibble +// shift/mask/bias-subtract/mask chain. +// +// For v in [0,15] the biased value is v-8. `v ^ 8` is exactly the 4-bit +// two's-complement pattern of v-8, because -8 == +8 (mod 16): +// v=0 -> 0x8 -> -8 v=7 -> 0xF -> -1 +// v=8 -> 0x0 -> 0 v=15 -> 0x7 -> +7 +// so the only remaining work is sign-extending bit 3 into bits 7:4 per byte. +// `sgn * 0x1E` does that with no cross-byte carry: 0x08 * 0x1E == 0xF0 exactly, +// and sgn is at most 0x08080808 so the product is at most 0xF0F0F0F0. +// +// A naive `nib - 0x08080808` would NOT work -- it borrows across byte lanes +// whenever a nibble is < 8. Shifts must be on uint, not int, so the >> is +// logical rather than arithmetic. +// +// ~5 ops per 4 weights vs ~22 for the per-nibble chain; bit-identical output. +uint widen_nibbles(const uint w, const uint parity) { + const uint nib = (parity == 0u) ? (w & 0x0F0F0F0Fu) : ((w >> 4u) & 0x0F0F0F0Fu); + const uint p = nib ^ 0x08080808u; + const uint sgn = p & 0x08080808u; + return p | (sgn * 0x1Eu); +} + +void main() { + const uvec2 tileID = uvec2(gl_WorkGroupID.xy); + const uvec2 warpInTile = uvec2( + gl_SubgroupID % SG_GRID_X, + gl_SubgroupID / SG_GRID_X); + + const uint K = uint(input_sizes.x); + const uint N = uint(output_sizes.x); + const uint N4 = (N + 3u) / 4u; + const uint nblocks_x_A = (K + 3u) >> 2u; + // A row stride in INT8 elements (dbuf4tr's binding is row-major int8, not + // the 4h4w ivec4 block layout dbuf4zpg used -- so A addressing below is in + // int8 elements, not int, and derived from nblocks_x_A so it matches the + // packer's `m_row * K4 + k4` addressing exactly (K % 4 == 0 makes them + // equal to K directly). + const uint a_row_stride_i8 = nblocks_x_A * 4u; + +#ifdef WEIGHT_INT4 + const uint num_groups = uint(num_groups_arg); + const uint CHUNKS_PER_GROUP = uint(K4_per_group) * 4u / WG_TILE_K; +#else + // Per-channel: a single quant "group" spanning all of K. The nested + // groups x chunks loop below collapses to a flat chunk loop, the wsum/wsc + // ping-pong never crosses a boundary, and the epilog runs exactly once. + const uint num_groups = 1u; + const uint CHUNKS_PER_GROUP = uint(num_groups_arg); +#endif + const uint num_chunks = num_groups * CHUNKS_PER_GROUP; + + const uint tile_m_start = WG_TILE_M * tileID.y; + const uint tile_n_start = WG_TILE_N * tileID.x; + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + result[i][j] = coopmat(0.0); + accum_int32[i][j] = coopmat(0); + } + } + + const uint K_BLOCKS_PER_CHUNK = WG_TILE_K >> 2u; + + // --- A staging tile map (dbuf4tr's, unmodified): one MMA_M x MMA_K coopmat + // tile per subgroup per slot, dealt round-robin across the + // NUM_SUBGROUPS subgroups so every subgroup participates. Replaces + // dbuf4zpg's per-thread (m4, k4) map / a_active guard entirely -- see + // design.md D3 for why this is reused as-is, not re-derived. + const uint A_TILES_M = WG_TILE_M / MMA_M; + const uint A_TILES_K = WG_TILE_K / MMA_K; // == NUM_K_SLABS + const uint NUM_A_TILES = A_TILES_M * A_TILES_K; + const uint A_TILES_PER_SG = (NUM_A_TILES + NUM_SUBGROUPS - 1u) / NUM_SUBGROUPS; + +#ifdef WEIGHT_INT4 + // --- B staging thread map: (block, col) slots; each slot extracts one + // ColumnMajor LDS uint (4 K-contiguous sign-extended int8) --- + const uint B_TOTAL_SLOTS = K_BLOCKS_PER_CHUNK * WG_TILE_N; + const uint B_SLOTS_PER_THREAD = B_TOTAL_SLOTS / WG_SIZE; + const uint N8_PER_TILE = WG_TILE_N >> 3u; +#else + // --- B staging thread map: one (k4, n4) ivec4 block per active thread --- + const uint B_FETCH_SLOTS = K_BLOCKS_PER_CHUNK * (WG_TILE_N >> 2u); + const uint N4_PER_TILE = WG_TILE_N >> 2u; + const uint b_k4_in_chunk = gl_LocalInvocationID.x / N4_PER_TILE; + const uint b_n_uint_col = gl_LocalInvocationID.x % N4_PER_TILE; + const bool b_active = gl_LocalInvocationID.x < B_FETCH_SLOTS; +#endif + + // ===== INTERVENTION F: hoist loop-invariant B staging index math ===== + // dbuf4zpg's, unchanged -- B staging is untouched by this file's A-staging + // swap. See dbuf4zpg's header for the full rationale (ablation-attributed + // -16.8% block, mostly index arithmetic not memory traffic). +#ifdef WEIGHT_INT4 + // bw3 ownership == dbuf4zpgbw2's (proven correct), one uvec2 element each. + uint b_v2off[V2_PER_THREAD]; // uvec2 element index within a slice + uint b_comp1[V2_PER_THREAD]; + uint b_par1[V2_PER_THREAD]; + uint b_n8b[V2_PER_THREAD]; + uint b_k4b[V2_PER_THREAD]; // k4 of the FIRST uint of the pair + [[unroll]] for (uint si = 0; si < V2_PER_THREAD; ++si) { + const uint e = gl_LocalInvocationID.x + si * WG_SIZE; // uvec2 index + const uint u0 = e * 2u; // first uint + const uint slab_idx = u0 / B_SLAB_U32; + const uint inslab = u0 % B_SLAB_U32; + const uint n_col = inslab / B_STRIDE_U32; + const uint k4_lo = inslab % B_STRIDE_U32; + const uint rem = n_col & 7u; + b_v2off[si] = e; + b_comp1[si] = rem & 3u; + b_par1[si] = rem >> 2u; + b_n8b[si] = (tile_n_start >> 3u) + (n_col >> 3u); + b_k4b[si] = slab_idx * (MMA_K >> 2u) + k4_lo; + } +#endif + + // Prefetch temp registers. temp_A is a coopmat array (dbuf4tr's A-staging + // technique); indices into it are [[unroll]]-resolved compile-time + // constants, never dynamic -- dynamic indexing of a coopmat array is + // exactly the construct the Xclipse/AMD-PAL compiler has miscompiled + // before. + coopmat + temp_A[A_TILES_PER_SG]; +#ifdef WEIGHT_INT4 + ivec4 temp_B[V2_PER_THREAD * 2u]; // bw3: 2 packed fetches per uvec2 + float temp_wsc; +#else + ivec4 temp_B; +#endif + + // ========================================================= + // PROLOGUE + // ========================================================= + if (gl_LocalInvocationID.x < (WG_TILE_M >> 2u)) { + const uint m4 = (tile_m_start >> 2u) + gl_LocalInvocationID.x; + const vec4 sc = vec4(texelFetch(t_int8_input_scales, ivec3(m4, 0, 0), 0)); + const ivec4 zp = texelFetch(t_int8_input_zps, ivec3(m4, 0, 0), 0); + const uint base = gl_LocalInvocationID.x * 4u; + ifs_sh[base + 0u] = sc.x; ifs_sh[base + 1u] = sc.y; + ifs_sh[base + 2u] = sc.z; ifs_sh[base + 3u] = sc.w; + izp_sh[base + 0u] = zp.x; izp_sh[base + 1u] = zp.y; + izp_sh[base + 2u] = zp.z; izp_sh[base + 3u] = zp.w; + } + // Group 0 weight scales -> slice 0, and the hoisted weight-side correction + // SUM_g wsc[g][n]*wsum[g][n] accumulated across ALL groups. dbuf4zpg's + // zp-hoist, unchanged. + if (gl_LocalInvocationID.x < WG_TILE_N) { + const uint n_idx = tile_n_start + gl_LocalInvocationID.x; + f16vec4 sv0 = t_weight_scales[n_idx >> 2u]; + wsc_sh[gl_LocalInvocationID.x] = float(sv0[n_idx & 3u]); + + float corr = 0.0; + for (uint g = 0; g < num_groups; ++g) { + f16vec4 sv = t_weight_scales[g * N4 + (n_idx >> 2u)]; + corr += float(sv[n_idx & 3u]) * float(t_weight_sums[g * N + n_idx]); + } + wcorr_sh[gl_LocalInvocationID.x] = corr; + } + memoryBarrierShared(); + barrier(); + + // NOTE: dbuf4zpg builds izp_bcast/ifs_bcast AFTER the group loop, not here + // -- that is the register-pressure saving zp-hoist buys. Unchanged. + + // dbuf4: prefetch chunk 0 into temp registers, THEN store to slice 0 (no + // barrier here -- the main loop's first iteration barriers before + // reading slice 0). + // + // A staging (dbuf4tr's technique): per-subgroup coopMatLoad straight from + // the row-major global buffer. + [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { + const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; + if (t < NUM_A_TILES) { + const uint tm = t / A_TILES_K; + const uint tk = t % A_TILES_K; + coopMatLoad( + temp_A[s], t_packed_int8_input, + (tile_m_start + tm * MMA_M) * a_row_stride_i8 + tk * MMA_K, + a_row_stride_i8, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#ifdef WEIGHT_INT4 + [[unroll]] for (uint si = 0; si < V2_PER_THREAD; ++si) { + [[unroll]] for (uint h = 0; h < 2u; ++h) { +#ifdef WEIGHT_BUFFER + temp_B[si * 2u + h] = t_packed_weight[(b_n8b[si] * nblocks_x_A) + b_k4b[si] + h]; +#else + temp_B[si * 2u + h] = texelFetch(t_packed_weight, ivec2(b_k4b[si] + h, b_n8b[si]), 0); +#endif + } + } +#else + if (b_active) { + const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; +#ifdef WEIGHT_BUFFER + temp_B = t_packed_weight[(b_k4_in_chunk * N4) + block_x_w]; +#else + temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, b_k4_in_chunk), 0); +#endif + } +#endif + { + // store chunk 0 -> slice 0 + // A staging (dbuf4tr's technique): coopMatStore into the same Ash_int8 + // slot layout dbuf4zpg's scalar scatter used to write. + [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { + const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; + if (t < NUM_A_TILES) { + const uint tm = t / A_TILES_K; + const uint tk = t % A_TILES_K; + coopMatStore( + temp_A[s], Ash_int8, + tk * A_SLAB_U32 + (tm * MMA_M) * A_STRIDE_U32, + A_STRIDE_U32, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#ifdef WEIGHT_INT4 + [[unroll]] for (uint si = 0; si < V2_PER_THREAD; ++si) { + // ONE ds_write_b64 per uvec2 element. + Bsh_v2[ b_v2off[si]] = uvec2( + widen_nibbles(uint(temp_B[si * 2u + 0u][b_comp1[si]]), b_par1[si]), + widen_nibbles(uint(temp_B[si * 2u + 1u][b_comp1[si]]), b_par1[si])); + } +#else + if (b_active) { + const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); + const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); + const uint n_col_base = b_n_uint_col * 4u; + [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { + // bw3: same address as (uvec2 element, component). This INT8 path + // writes 4 different columns for one k4, so no wide store applies. + { + const uint u = slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab; + Bsh_v2[u / 2u][u % 2u] = uint(temp_B[n_in_blk]); + } + } + } +#endif + } + + // ========================================================= + // MAIN LOOP (dbuf4) — nested groups x chunks (kept nested; flattening it + // with a conditional coopmat epilog crashes the Xclipse PAL compiler at + // large spec-resolved trip counts). One barrier per chunk. Chunk + // iteration (global index `chunk`): + // 1. barrier — A/B slice (chunk%2) fully written; on the first chunk + // of group g, wsc slice (g%2) is too. + // 2. prefetch — chunk+1 (A tiles, B blocks) into temp; when chunk+1 + // starts a new group, also its wsc element. Skipped + // entirely on the final chunk. + // 3. int8 MMA — on slice (chunk%2) into accum_int32. + // 4. store — temp -> A/B slice ((chunk+1)%2), unpacking the weight; + // on a group boundary, wsc -> slice ((g+1)%2). + // The group epilog runs unconditionally at the tail of each group. + // ========================================================= + uint chunk = 0; + for (uint group_i = 0; group_i < num_groups; ++group_i) { + for (uint inner = 0; inner < CHUNKS_PER_GROUP; ++inner, ++chunk) { + const bool has_next = chunk + 1u < num_chunks; + const bool group_crossing = has_next && (inner + 1u == CHUNKS_PER_GROUP); + const uint cur_a = (chunk % 2u) * ASH_SLICE_U32; + const uint cur_b_v2 = (chunk % 2u) * BSH_SLICE_V2; + const uint nxt_a = ((chunk + 1u) % 2u) * ASH_SLICE_U32; + const uint nxt_b_v2 = ((chunk + 1u) % 2u) * BSH_SLICE_V2; + + // coopmat-lds-fence: barrier() alone does NOT order shared stores against a + // subsequent coopMatLoad on the M51 Xclipse/AMD-PAL driver -- symptom is one + // stale MMA_M-row band of A, all columns, ~2.5% of runs, no crash (observed + // 2026-09-02 in sdpa_compute_out_coopmat.glsl). Measured cost of the fence: + // none (see this change's results). See memory + // `coopmat-lds-needs-explicit-memorybarriershared`. + memoryBarrierShared(); + barrier(); + + // --- 2. prefetch chunk+1 -> temp --- + if (has_next) { + const uint chunkK_nxt = (chunk + 1u) * WG_TILE_K; + // A staging (dbuf4tr's technique): coopMatLoad straight from global. + [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { + const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; + if (t < NUM_A_TILES) { + const uint tm = t / A_TILES_K; + const uint tk = t % A_TILES_K; + coopMatLoad( + temp_A[s], t_packed_int8_input, + (tile_m_start + tm * MMA_M) * a_row_stride_i8 + chunkK_nxt + + tk * MMA_K, + a_row_stride_i8, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#ifdef WEIGHT_INT4 + [[unroll]] for (uint si = 0; si < V2_PER_THREAD; ++si) { + [[unroll]] for (uint h = 0; h < 2u; ++h) { + const uint k4_blk = (chunkK_nxt >> 2u) + b_k4b[si] + h; +#ifdef WEIGHT_BUFFER + temp_B[si * 2u + h] = t_packed_weight[(b_n8b[si] * nblocks_x_A) + k4_blk]; +#else + temp_B[si * 2u + h] = texelFetch(t_packed_weight, ivec2(k4_blk, b_n8b[si]), 0); +#endif + } + } + if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { + const uint n_idx = tile_n_start + gl_LocalInvocationID.x; + f16vec4 sv = t_weight_scales[(group_i + 1u) * N4 + (n_idx >> 2u)]; + temp_wsc = float(sv[n_idx & 3u]); + } +#else + if (b_active) { + const uint block_y_w = (chunkK_nxt >> 2u) + b_k4_in_chunk; + const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; +#ifdef WEIGHT_BUFFER + temp_B = t_packed_weight[(block_y_w * N4) + block_x_w]; +#else + temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, block_y_w), 0); +#endif + } +#endif + } + + // --- 3. int8 MMA on the cur slice --- + [[unroll]] for (uint k = 0; k < NUM_K_SLABS; ++k) { + const uint slab_a_base_u32 = cur_a + k * A_SLAB_U32; + const uint slab_b_base_v2 = cur_b_v2 + k * B_SLAB_V2; + + coopmat matA[MMAS_PER_SG_M]; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + const uint row_a = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + coopMatLoad( + matA[i], Ash_int8, + slab_a_base_u32 + row_a * A_STRIDE_U32, + A_STRIDE_U32, + gl_CooperativeMatrixLayoutRowMajor); + } + + coopmat matB; + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint col_b = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopMatLoad( + matB, Bsh_v2, + slab_b_base_v2 + col_b * B_STRIDE_V2, + B_STRIDE_V2, + gl_CooperativeMatrixLayoutColumnMajor); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + accum_int32[i][j] = coopMatMulAdd(matA[i], matB, accum_int32[i][j]); + } + } + } + + // --- 4. store temp (chunk+1) -> nxt slice --- + if (has_next) { + // A staging (dbuf4tr's technique): coopMatStore into the nxt slice. + [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { + const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; + if (t < NUM_A_TILES) { + const uint tm = t / A_TILES_K; + const uint tk = t % A_TILES_K; + coopMatStore( + temp_A[s], Ash_int8, + nxt_a + tk * A_SLAB_U32 + (tm * MMA_M) * A_STRIDE_U32, + A_STRIDE_U32, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#ifdef WEIGHT_INT4 + [[unroll]] for (uint si = 0; si < V2_PER_THREAD; ++si) { + // ONE ds_write_b64 per uvec2 element. + Bsh_v2[nxt_b_v2 + b_v2off[si]] = uvec2( + widen_nibbles(uint(temp_B[si * 2u + 0u][b_comp1[si]]), b_par1[si]), + widen_nibbles(uint(temp_B[si * 2u + 1u][b_comp1[si]]), b_par1[si])); + } + if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { + const uint wbase_nxt = ((group_i + 1u) % 2u) * WG_TILE_N; + wsc_sh[wbase_nxt + gl_LocalInvocationID.x] = temp_wsc; + } +#else + if (b_active) { + const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); + const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); + const uint n_col_base = b_n_uint_col * 4u; + [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { + { + const uint u = slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab; + Bsh_v2[nxt_b_v2 + u / 2u][u % 2u] = uint(temp_B[n_in_blk]); + } + } + } +#endif + } + } // chunks + + // --- Group epilog: scale-only accumulate, reset accum --- + // dbuf4zpg's, unchanged. Just result += float(acc) * wsc. The + // zero-point subtract and the ifs multiply are hoisted out of the group + // loop (applied once below). + { + const uint wbase = (group_i % 2u) * WG_TILE_N; + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + + coopmat wsc_bcast; + coopMatLoad( + wsc_bcast, wsc_sh, + wbase + local_n_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutRowMajor); + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + result[i][j] += + coopmat( + accum_int32[i][j]) * wsc_bcast; + accum_int32[i][j] = coopmat(0); + } + } + } + } // groups + + // --- Hoisted correction, applied ONCE: -------------------------------- + // result = ifs * ( result - izp * SUM_g wsc_g*wsum_g ) + // dbuf4zpg's, unchanged. izp/ifs are loaded here rather than before the + // group loop so they are not live across it. + { + coopmat + izpf_bcast[MMAS_PER_SG_M]; + coopmat + ifs_bcast[MMAS_PER_SG_M]; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + const uint local_m_base = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + coopmat izp_i; + coopMatLoad( + izp_i, izp_sh, local_m_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutColumnMajor); + izpf_bcast[i] = + coopmat(izp_i); + coopMatLoad( + ifs_bcast[i], ifs_sh, local_m_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutColumnMajor); + } + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopmat wcorr_bcast; + coopMatLoad( + wcorr_bcast, wcorr_sh, local_n_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutRowMajor); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + result[i][j] = ifs_bcast[i] * (result[i][j] - izpf_bcast[i] * wcorr_bcast); + } + } + } + + // --- Bias (optional) --- +#ifdef HAS_BIAS + if (apply_bias > 0) { + for (uint t = gl_LocalInvocationID.x; t < WG_TILE_N; t += WG_SIZE) { + bias_sh[t] = float(t_bias[tile_n_start + t]); + } + memoryBarrierShared(); + barrier(); + } +#endif + + // --- Store result tile --- + // N for the store address math MUST come from the spec constant, not the + // sizes UBO (see out_N_arg above). +#ifdef IO_TEXTURE + // Epilogue iteration i drains accumulator row-block i from EVERY subgroup + // into Csh_out at once, so the SG_GRID_Y bands it holds are disjoint global + // row ranges; the whole workgroup then imageStores them. lr / MMA_M is the + // writing subgroup's warpInTile.y, so the global row reproduces the buffer + // path's gi exactly. + // + // PORTABILITY NOTE: the barrier() in the loop body keeps this loop rolled + // despite [[unroll]], so result[i][j] IS dynamically indexed. Coopmat arrays + // are opaque per-lane storage and dynamic indexing is exactly the construct + // the Xclipse/AMD-PAL compiler has broken before -- check this first if the + // texture variants miscompile on M51. + const uint CSH_TEXELS_PER_ROW = WG_TILE_N / 4u; + const uint CSH_TEXELS = CSH_ROWS * CSH_TEXELS_PER_ROW; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + // Guards Csh_out against the previous iteration's readers. Inert on i == 0 + // but must stay unconditional to remain workgroup-uniform. + // coopmat-lds-fence: barrier() alone does NOT order shared stores against a + // subsequent coopMatLoad on the M51 Xclipse/AMD-PAL driver -- symptom is one + // stale MMA_M-row band of A, all columns, ~2.5% of runs, no crash (observed + // 2026-09-02 in sdpa_compute_out_coopmat.glsl). Measured cost of the fence: + // none (see this change's results). See memory + // `coopmat-lds-needs-explicit-memorybarriershared`. + memoryBarrierShared(); + barrier(); + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { +#ifdef HAS_BIAS + if (apply_bias > 0) { + const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopmat bias_tile; + coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); + result[i][j] += bias_tile; + } +#endif + coopmat out_tile = + coopmat(result[i][j]); + coopMatStore( + out_tile, Csh_out, + warpInTile.y * MMA_M * WG_TILE_N + + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j), + WG_TILE_N, + gl_CooperativeMatrixLayoutRowMajor); + } + memoryBarrierShared(); + barrier(); + + for (uint t = gl_LocalInvocationID.x; t < CSH_TEXELS; t += WG_SIZE) { + const uint lr = t / CSH_TEXELS_PER_ROW; + const uint lc4 = t % CSH_TEXELS_PER_ROW; + const uint m = + tile_m_start + (lr / MMA_M) * SG_TILE_M + i * MMA_M + (lr % MMA_M); + const uint base = lr * WG_TILE_N + lc4 * 4u; + imageStore( + t_output, + ivec3(tile_n_start / 4u + lc4, m, 0), + vec4( + float(Csh_out[base]), + float(Csh_out[base + 1u]), + float(Csh_out[base + 2u]), + float(Csh_out[base + 3u]))); + } + } +#else + const uint N_out = uint(out_N_arg); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint gi = tile_m_start + MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + const uint gj = tile_n_start + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + +#ifdef HAS_BIAS + if (apply_bias > 0) { + const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopmat bias_tile; + coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); + result[i][j] += bias_tile; + } +#endif + + coopmat out_tile = + coopmat(result[i][j]); + coopMatStore( + out_tile, t_output, + gi * N_out + gj, N_out, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#endif // IO_TEXTURE +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw3.yaml b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw3.yaml new file mode 100644 index 00000000000..d167ca95809 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw3.yaml @@ -0,0 +1,60 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# "zpg" + "-tr" combination: tsweep_dbuf4zpg with its per-thread scalar +# A-staging replaced by tsweep_dbuf4tr's coopMat-mediated A staging. B +# staging, zp-hoist, byte-parallel nibble widening are all dbuf4zpg's, +# unchanged -- only A-staging differs. Requires t_packed_int8_input in the +# ROW-MAJOR kPackedInt8_4W layout (same requirement as tsweep_dbuf4tr). +# +# PROMOTED 2026-09-01 as the shipped default -- see QuantizedLinear.cpp's +# dq8ca_coopmat_variant() for the full validation record (10/10 buffer, 6/6 +# texture3d correctness across 1B/3B/8B; +4.2% real e2e prefill; 46.50% -> +# 49.94% efficiency on 8B). Also selectable explicitly via +# ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4zpgbw3_txkgs. +# See openspec/changes/archive/2026-08-31-dq8ca-tr-staged-a-on-zpg. +# +# Single seed tile: the prior dbuf4zpg default's own tile (t128x64k32g42s32). +# A re-sweep against this shader's own (lower) register-pressure profile was +# run as a follow-up (coopmat-tr-tilesweep-4w-port) and found no better tile +# -- this remains the best known geometry. + +linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw3: + parameter_names_with_default_values: + PRECISION: highp + HAS_BIAS: false + IO_STORAGE: buffer + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + MMA_M: 16 + MMA_N: 16 + MMA_K: 16 + WG_TILE_M: 64 + WG_TILE_N: 32 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + shader_variants: + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw3_t128x64k32g42s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw3_t128x64k32g42s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbwr.glsl b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbwr.glsl new file mode 100644 index 00000000000..e187ae60d97 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbwr.glsl @@ -0,0 +1,783 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +/* + * "zpg" + "-tr" combination: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpg.glsl + * with its per-thread scalar A-staging replaced by + * linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr.glsl's coopMat-mediated A + * staging, PLUS "bw": B staging ownership moved from (block,col) slots spread + * over 4 threads to ONE thread per (K-slab, column), writing its 4 uints at + * consecutive LDS indices so they are merge-eligible. Shared-array types, the + * compute-side loads, the quantization math and every safeguard are + * dbuf4zpgtr's, unchanged. This is + * an ADDITIVE combination, not a redesign: every non-A-staging block below + * (B staging: coalesced write, no skew; zp-hoist: izp/ifs applied once after + * the group loop via wcorr_sh; byte-parallel nibble widening; static + * A_ALWAYS_ACTIVE branch elision -- N/A here, see below; group epilog; + * bias/store epilogue) is byte-identical to dbuf4zpg's. Only the A-staging + * block (prologue load+store, main-loop prefetch+store) is dbuf4tr's, + * verbatim. + * + * dbuf4zpg's per-thread A staging used an `a_active` guard (statically always + * true when A_ACTIVE_THREADS == WG_SIZE, via the A_MAP_FULL-gated + * A_ALWAYS_ACTIVE macro). dbuf4tr's per-SUBGROUP tile map has no equivalent + * concept -- every subgroup participates via a `t < NUM_A_TILES` guard that + * depends only on gl_SubgroupID, not gl_LocalInvocationID.x -- so + * A_MAP_FULL/A_ALWAYS_ACTIVE is dropped entirely in this file; it would be + * dead code for the new A-staging block. + * + * Rationale for combining this way (not the reverse) and why this is worth + * building at all: see this change's design.md D0-D3. In short -- the only + * existing measurement of dbuf4tr's A-staging technique (28.72-30.51%, + * dq8ca-arch-redesign) was taken against dbuf4tr's own pre-zpg baseline + * (old B skew, no byte-parallel widening, no branch elision) -- a materially + * weaker shader than the 46.49-46.50% dbuf4zpg this file now combines it + * with. This file exists to answer whether that combination performs + * differently now that register pressure is already reduced. + * + * A staging (the actual delta from dbuf4zpg): + * dbuf4zpg: per-thread (m4, k4) ivec4 fetch, hoisted a_lds_off0/a_glb_row; + * only A_ACTIVE_THREADS invocations participate, each scattering + * 4 rows into Ash_int8 with 4 scalar stores. + * this file: per-SUBGROUP MMA_M x MMA_K tile fetch via coopMatLoad straight + * from a ROW-MAJOR (kPackedInt8_4W) int8 activation buffer, then + * coopMatStore into the same Ash_int8 slot -- dbuf4tr's mapping, + * unmodified (not re-derived; see design.md D3). + * + * t_packed_int8_input is therefore bound the same way dbuf4tr binds it: a + * SCALAR int8_t array in the kPackedInt8_4W layout (plain row-major int8, + * row stride K), produced by quantize_and_pack_4w_with_group_sums.glsl. + * QuantizedLinear.cpp's dq8ca_variant_wants_rowmajor_a() must recognize this + * file's variant token (tsweep_dbuf4zpgtr_t...) the same way it already + * recognizes tsweep_dbuf4tr_t/trm_t/trd_t, so graph-build time (packer + * selection) and dispatch time (kernel selection) cannot disagree. + * + * B CANNOT be coopmat-staged (int4 nibble unpack; a coopmat's per-lane layout + * is opaque to hand-assembly from unpacked registers) -- unchanged from both + * parent files. B staging below is dbuf4zpg's byte-parallel, coalesced, + * no-skew version, untouched. + * + * The loop structure is dbuf4's (both parents share it), unchanged: + * prologue: prefetch chunk 0 -> temp, store to slice 0 (no barrier) + * per iter: barrier -> prefetch(next) -> MMA(cur) -> store(next) + * kept nested (groups x chunks) with an unconditional group epilog -- + * flattening it crashes the Xclipse PAL compiler at large spec-resolved trip + * counts (see dbuf2's own header). + * + * Selected via + * ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4zpgtr_txkgs<32|64> + * (QuantizedLinear.cpp), additive to the tsweep_dbuf4zpg_t..., tsweep_dbuf4tr_t... + * and tsweep_t... namespaces. NOT the default -- unvalidated until it passes + * repeated test_llama_microbench --correctness-only runs (see + * dq8ca_coopmat_variant()'s comment on why a single pass is not proof). + * + * Performs: out[M,N] = dequant(int8_act) * dequant(int_w) (+ bias) + * via coopmat x coopmat -> coopmat on the matrix unit. + * + * Hard preconditions (dbuf4zpg's, plus dbuf4tr's row-major/alignment ones): + * M % WG_TILE_M == 0, N % WG_TILE_N == 0, K % WG_TILE_K == 0, + * group_size % WG_TILE_K == 0, K % 4 == 0, + * WG_TILE_M % MMA_M == 0, WG_TILE_K % MMA_K == 0, + * t_packed_int8_input in kPackedInt8_4W (row-major) layout, + * device exposes coopmatx-> at 16x16x16. + */ + +#version 450 core + +#extension GL_KHR_cooperative_matrix : require +#extension GL_KHR_memory_scope_semantics : require +#extension GL_KHR_shader_subgroup_basic : enable +#extension GL_EXT_shader_explicit_arithmetic_types : require +#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require +// 8-bit SSBO access: A is bound as a scalar int8_t array so that the +// coopMatLoad below has a MATCHING component type (see dbuf4tr's header for +// why the type must match on this driver). +#extension GL_EXT_shader_8bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_EXT_control_flow_attributes : enable + +#define PRECISION ${PRECISION} + +$if WEIGHT_NBITS == 4: + #define WEIGHT_INT4 + +$if HAS_BIAS: + #define HAS_BIAS + +$if WEIGHT_STORAGE == "buffer": + #define WEIGHT_BUFFER + +$if IO_STORAGE == "texture3d": + #define IO_TEXTURE + +layout(std430) buffer; + +#include "common.glslh" + +// Bindings — match add_linear_dqa_qw_node arg order: +// output(0), fp_input(1), packed_int8_input(2), int_input_sums(3 - unused), +// input_scales(4), input_zps(5), packed_weight(6), weight_sums(7), +// weight_scales(8), bias(9). +${layout_declare_tensor(B, "w", "t_output", "half", IO_STORAGE, is_scalar_array=True)} +// t_input is unread here -- the activations arrive already quantized in +// t_packed_int8_input -- but stays declared so the binding layout matches the +// dispatch site. It tracks IO_STORAGE so the two IO tensors stay consistent. +${layout_declare_tensor(B, "r", "t_input", "half", IO_STORAGE, is_scalar_array=False)} +// ROW-MAJOR (kPackedInt8_4W) packed activations, bound as a scalar int8_t +// array (row stride = K int8) -- dbuf4tr's binding, unchanged. The stock +// 4h4w layout dbuf4zpg uses is NOT row-major (component index selects a row, +// non-affine), so it cannot be addressed by any coopMatLoad. +${layout_declare_tensor(B, "r", "t_packed_int8_input", "int8", "buffer", is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_int8_input_sums", "int", "buffer", is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_int8_input_scales", "half", "texture3d")} +${layout_declare_tensor(B, "r", "t_int8_input_zps", "int8", "texture3d")} +${layout_declare_tensor(B, "r", "t_packed_weight", "int", WEIGHT_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_weight_sums", "int", "buffer", is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_weight_scales", "half", "buffer", is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_bias", "half", "buffer", is_scalar_array=True)} + +${layout_declare_ubo(B, "ivec4", "output_sizes")} +${layout_declare_ubo(B, "ivec4", "input_sizes")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +${layout_declare_spec_const(C, "int", "apply_bias", "0")} +// INT4 only; inert (0) for INT8 so the dispatcher's spec list lines up. +${layout_declare_spec_const(C, "int", "K4_per_group", "0")} +${layout_declare_spec_const(C, "int", "num_groups_arg", "0")} +${layout_declare_spec_const(C, "int", "out_N_arg", "0")} + +// Tile geometry +const uint MMA_M = ${MMA_M}; +const uint MMA_N = ${MMA_N}; +const uint MMA_K = ${MMA_K}; + +const uint WG_TILE_M = ${WG_TILE_M}; +const uint WG_TILE_N = ${WG_TILE_N}; +const uint WG_TILE_K = ${WG_TILE_K}; + +const uint SG_GRID_X = ${SG_GRID_X}; +const uint SG_GRID_Y = ${SG_GRID_Y}; +const uint SUBGROUP_SIZE = ${SUBGROUP_SIZE}; +const uint NUM_SUBGROUPS = SG_GRID_X * SG_GRID_Y; +const uint WG_SIZE = NUM_SUBGROUPS * SUBGROUP_SIZE; + +const uint SG_TILE_M = WG_TILE_M / SG_GRID_Y; +const uint SG_TILE_N = WG_TILE_N / SG_GRID_X; +const uint MMAS_PER_SG_M = SG_TILE_M / MMA_M; +const uint MMAS_PER_SG_N = SG_TILE_N / MMA_N; + +const uint A_SLAB_INT8 = WG_TILE_M * MMA_K; +const uint B_USEFUL_U32 = MMA_K / 4u; +// No skew + coalesced write -- dbuf4zpg's B fix, unchanged (this file does +// not touch B staging at all). +const uint B_STRIDE_U32 = B_USEFUL_U32; +const uint B_SLAB_U32 = WG_TILE_N * B_STRIDE_U32; +const uint NUM_K_SLABS = WG_TILE_K / MMA_K; + +const uint A_SLAB_U32 = A_SLAB_INT8 / 4u; +const uint A_STRIDE_U32 = MMA_K / 4u; + +// One ping-pong slice covers all K-slabs of one chunk. +const uint ASH_SLICE_U32 = NUM_K_SLABS * A_SLAB_U32; +const uint BSH_SLICE_U32 = NUM_K_SLABS * B_SLAB_U32; + +// Double-buffered MMA operand staging. +shared uint Ash_int8[2u * ASH_SLICE_U32]; +shared uint Bsh_int8[2u * BSH_SLICE_U32]; + +// Per-WG-tile-row activation params (loaded ONCE at WG start; constant +// across groups). +shared int izp_sh[WG_TILE_M]; // int32 (cast from int8 source) for broadcast +shared float ifs_sh[WG_TILE_M]; // float32 (cast from fp16 source) for broadcast + +// Per-(group, output-channel) weight params, ping-ponged by group parity. +// (For per-channel INT8 only slice 0 is ever used.) +shared float wsc_sh[2u * WG_TILE_N]; +// SUM_g wsc[g][n]*wsum[g][n] per output channel -- weight-side only, so it is +// accumulated once in the prologue. dbuf4zpg's zp-hoist, unchanged. +shared float wcorr_sh[WG_TILE_N]; + +#ifdef HAS_BIAS +shared float bias_sh[WG_TILE_N]; +#endif + +#ifdef IO_TEXTURE +// Result staging for the imageStore epilogue, mirroring the fp16 kernel: +// SG_GRID_Y bands of MMA_M rows, each WG_TILE_N wide, row-major. A full +// WG_TILE_M x WG_TILE_N buffer would cost SG_GRID_Y/MMAS_PER_SG_M x more LDS +// and wreck occupancy. float16_t-typed because coopMatStore needs it. +const uint CSH_ROWS = SG_GRID_Y * MMA_M; +shared float16_t Csh_out[CSH_ROWS * WG_TILE_N]; +#endif + +// Running fp32 accumulator (across all groups). +coopmat + result[MMAS_PER_SG_M][MMAS_PER_SG_N]; + +// Per-group int32 MMA accumulator. +coopmat + accum_int32[MMAS_PER_SG_M][MMAS_PER_SG_N]; + + +// Byte-parallel int4 -> int8 widening. dbuf4zpg's, unchanged (B-side only). +// +// The four nibbles this shader needs from one packed uint are ALREADY one per +// byte (bits 3:0 of each byte for parity 0, bits 7:4 for parity 1), so all four +// can be widened at once instead of with a per-nibble +// shift/mask/bias-subtract/mask chain. +// +// For v in [0,15] the biased value is v-8. `v ^ 8` is exactly the 4-bit +// two's-complement pattern of v-8, because -8 == +8 (mod 16): +// v=0 -> 0x8 -> -8 v=7 -> 0xF -> -1 +// v=8 -> 0x0 -> 0 v=15 -> 0x7 -> +7 +// so the only remaining work is sign-extending bit 3 into bits 7:4 per byte. +// `sgn * 0x1E` does that with no cross-byte carry: 0x08 * 0x1E == 0xF0 exactly, +// and sgn is at most 0x08080808 so the product is at most 0xF0F0F0F0. +// +// A naive `nib - 0x08080808` would NOT work -- it borrows across byte lanes +// whenever a nibble is < 8. Shifts must be on uint, not int, so the >> is +// logical rather than arithmetic. +// +// ~5 ops per 4 weights vs ~22 for the per-nibble chain; bit-identical output. +uint widen_nibbles(const uint w, const uint parity) { + const uint nib = (parity == 0u) ? (w & 0x0F0F0F0Fu) : ((w >> 4u) & 0x0F0F0F0Fu); + const uint p = nib ^ 0x08080808u; + const uint sgn = p & 0x08080808u; + return p | (sgn * 0x1Eu); +} + +void main() { + const uvec2 tileID = uvec2(gl_WorkGroupID.xy); + const uvec2 warpInTile = uvec2( + gl_SubgroupID % SG_GRID_X, + gl_SubgroupID / SG_GRID_X); + + const uint K = uint(input_sizes.x); + const uint N = uint(output_sizes.x); + const uint N4 = (N + 3u) / 4u; + const uint nblocks_x_A = (K + 3u) >> 2u; + // A row stride in INT8 elements (dbuf4tr's binding is row-major int8, not + // the 4h4w ivec4 block layout dbuf4zpg used -- so A addressing below is in + // int8 elements, not int, and derived from nblocks_x_A so it matches the + // packer's `m_row * K4 + k4` addressing exactly (K % 4 == 0 makes them + // equal to K directly). + const uint a_row_stride_i8 = nblocks_x_A * 4u; + +#ifdef WEIGHT_INT4 + const uint num_groups = uint(num_groups_arg); + const uint CHUNKS_PER_GROUP = uint(K4_per_group) * 4u / WG_TILE_K; +#else + // Per-channel: a single quant "group" spanning all of K. The nested + // groups x chunks loop below collapses to a flat chunk loop, the wsum/wsc + // ping-pong never crosses a boundary, and the epilog runs exactly once. + const uint num_groups = 1u; + const uint CHUNKS_PER_GROUP = uint(num_groups_arg); +#endif + const uint num_chunks = num_groups * CHUNKS_PER_GROUP; + + const uint tile_m_start = WG_TILE_M * tileID.y; + const uint tile_n_start = WG_TILE_N * tileID.x; + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + result[i][j] = coopmat(0.0); + accum_int32[i][j] = coopmat(0); + } + } + + const uint K_BLOCKS_PER_CHUNK = WG_TILE_K >> 2u; + + // --- A staging tile map (dbuf4tr's, unmodified): one MMA_M x MMA_K coopmat + // tile per subgroup per slot, dealt round-robin across the + // NUM_SUBGROUPS subgroups so every subgroup participates. Replaces + // dbuf4zpg's per-thread (m4, k4) map / a_active guard entirely -- see + // design.md D3 for why this is reused as-is, not re-derived. + const uint A_TILES_M = WG_TILE_M / MMA_M; + const uint A_TILES_K = WG_TILE_K / MMA_K; // == NUM_K_SLABS + const uint NUM_A_TILES = A_TILES_M * A_TILES_K; + const uint A_TILES_PER_SG = (NUM_A_TILES + NUM_SUBGROUPS - 1u) / NUM_SUBGROUPS; + +#ifdef WEIGHT_INT4 + // --- B staging thread map: (block, col) slots; each slot extracts one + // ColumnMajor LDS uint (4 K-contiguous sign-extended int8) --- + const uint B_TOTAL_SLOTS = K_BLOCKS_PER_CHUNK * WG_TILE_N; + const uint B_SLOTS_PER_THREAD = B_TOTAL_SLOTS / WG_SIZE; + const uint N8_PER_TILE = WG_TILE_N >> 3u; +#else + // --- B staging thread map: one (k4, n4) ivec4 block per active thread --- + const uint B_FETCH_SLOTS = K_BLOCKS_PER_CHUNK * (WG_TILE_N >> 2u); + const uint N4_PER_TILE = WG_TILE_N >> 2u; + const uint b_k4_in_chunk = gl_LocalInvocationID.x / N4_PER_TILE; + const uint b_n_uint_col = gl_LocalInvocationID.x % N4_PER_TILE; + const bool b_active = gl_LocalInvocationID.x < B_FETCH_SLOTS; +#endif + + // ===== INTERVENTION F: hoist loop-invariant B staging index math ===== + // dbuf4zpg's, unchanged -- B staging is untouched by this file's A-staging + // swap. See dbuf4zpg's header for the full rationale (ablation-attributed + // -16.8% block, mostly index arithmetic not memory traffic). +#ifdef WEIGHT_INT4 + // ===== bw: SOURCE-side change only -- one thread owns a whole (slab, col) + // column-slab and writes its 4 uints at CONSECUTIVE indices, so the four + // ds_write_b32 are merge-eligible. The shared array stays `uint` (the proven + // 4x packing ratio) and the compute-side coopMatLoad is byte-for-byte + // unchanged -- deliberately NOT the uvec4 retype, which fails correctness on + // this driver (16x packing; see the 8x coopMatStore precedent). + // comp/par depend only on the column, so a thread's 4 fetches share them and + // differ only in k4: 4 consecutive k4 blocks of the same n8 texel column. + const uint B_OWNERS = NUM_K_SLABS * WG_TILE_N; + const bool b_owns = gl_LocalInvocationID.x < B_OWNERS; + const uint b_slab = gl_LocalInvocationID.x / WG_TILE_N; + const uint b_col = gl_LocalInvocationID.x % WG_TILE_N; + const uint b_rem = b_col & 7u; + const uint b_comp1 = b_rem & 3u; + const uint b_par1 = b_rem >> 2u; + const uint b_n8blk1 = (tile_n_start >> 3u) + (b_col >> 3u); + const uint b_k4base = b_slab * (MMA_K >> 2u); + // same address the destination-oriented map produced, just all 4 in one thread + const uint b_lds_base = b_slab * B_SLAB_U32 + b_col * B_STRIDE_U32; +#endif + + // Prefetch temp registers. temp_A is a coopmat array (dbuf4tr's A-staging + // technique); indices into it are [[unroll]]-resolved compile-time + // constants, never dynamic -- dynamic indexing of a coopmat array is + // exactly the construct the Xclipse/AMD-PAL compiler has miscompiled + // before. + coopmat + temp_A[A_TILES_PER_SG]; +#ifdef WEIGHT_INT4 + ivec4 temp_B[MMA_K >> 2u]; // bw: one per k4 of the owned column-slab + float temp_wsc; +#else + ivec4 temp_B; +#endif + + // ========================================================= + // PROLOGUE + // ========================================================= + if (gl_LocalInvocationID.x < (WG_TILE_M >> 2u)) { + const uint m4 = (tile_m_start >> 2u) + gl_LocalInvocationID.x; + const vec4 sc = vec4(texelFetch(t_int8_input_scales, ivec3(m4, 0, 0), 0)); + const ivec4 zp = texelFetch(t_int8_input_zps, ivec3(m4, 0, 0), 0); + const uint base = gl_LocalInvocationID.x * 4u; + ifs_sh[base + 0u] = sc.x; ifs_sh[base + 1u] = sc.y; + ifs_sh[base + 2u] = sc.z; ifs_sh[base + 3u] = sc.w; + izp_sh[base + 0u] = zp.x; izp_sh[base + 1u] = zp.y; + izp_sh[base + 2u] = zp.z; izp_sh[base + 3u] = zp.w; + } + // Group 0 weight scales -> slice 0, and the hoisted weight-side correction + // SUM_g wsc[g][n]*wsum[g][n] accumulated across ALL groups. dbuf4zpg's + // zp-hoist, unchanged. + if (gl_LocalInvocationID.x < WG_TILE_N) { + const uint n_idx = tile_n_start + gl_LocalInvocationID.x; + f16vec4 sv0 = t_weight_scales[n_idx >> 2u]; + wsc_sh[gl_LocalInvocationID.x] = float(sv0[n_idx & 3u]); + + float corr = 0.0; + for (uint g = 0; g < num_groups; ++g) { + f16vec4 sv = t_weight_scales[g * N4 + (n_idx >> 2u)]; + corr += float(sv[n_idx & 3u]) * float(t_weight_sums[g * N + n_idx]); + } + wcorr_sh[gl_LocalInvocationID.x] = corr; + } + memoryBarrierShared(); + barrier(); + + // NOTE: dbuf4zpg builds izp_bcast/ifs_bcast AFTER the group loop, not here + // -- that is the register-pressure saving zp-hoist buys. Unchanged. + + // dbuf4: prefetch chunk 0 into temp registers, THEN store to slice 0 (no + // barrier here -- the main loop's first iteration barriers before + // reading slice 0). + // + // A staging (dbuf4tr's technique): per-subgroup coopMatLoad straight from + // the row-major global buffer. + [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { + const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; + if (t < NUM_A_TILES) { + const uint tm = t / A_TILES_K; + const uint tk = t % A_TILES_K; + coopMatLoad( + temp_A[s], t_packed_int8_input, + (tile_m_start + tm * MMA_M) * a_row_stride_i8 + tk * MMA_K, + a_row_stride_i8, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#ifdef WEIGHT_INT4 + if (b_owns) { + [[unroll]] for (uint kk = 0; kk < (MMA_K >> 2u); ++kk) { +#ifdef WEIGHT_BUFFER + temp_B[kk] = t_packed_weight[(b_n8blk1 * nblocks_x_A) + b_k4base + kk]; +#else + temp_B[kk] = texelFetch(t_packed_weight, ivec2(b_k4base + kk, b_n8blk1), 0); +#endif + } + } +#else + if (b_active) { + const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; +#ifdef WEIGHT_BUFFER + temp_B = t_packed_weight[(b_k4_in_chunk * N4) + block_x_w]; +#else + temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, b_k4_in_chunk), 0); +#endif + } +#endif + { + // store chunk 0 -> slice 0 + // A staging (dbuf4tr's technique): coopMatStore into the same Ash_int8 + // slot layout dbuf4zpg's scalar scatter used to write. + [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { + const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; + if (t < NUM_A_TILES) { + const uint tm = t / A_TILES_K; + const uint tk = t % A_TILES_K; + coopMatStore( + temp_A[s], Ash_int8, + tk * A_SLAB_U32 + (tm * MMA_M) * A_STRIDE_U32, + A_STRIDE_U32, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#ifdef WEIGHT_INT4 + if (b_owns) { + // 4 CONSECUTIVE uints from one thread -> merge-eligible into a wider + // store, versus 4 scalar stores from 4 different threads before. + [[unroll]] for (uint q = 0; q < (MMA_K >> 2u); ++q) { + const uint kk = (MMA_K >> 2u) - 1u - q; // DIAGNOSTIC: reverse order + Bsh_int8[b_lds_base + kk] = + widen_nibbles(uint(temp_B[kk][b_comp1]), b_par1); + } + } +#else + if (b_active) { + const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); + const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); + const uint n_col_base = b_n_uint_col * 4u; + [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { + Bsh_int8[slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab] = + uint(temp_B[n_in_blk]); + } + } +#endif + } + + // ========================================================= + // MAIN LOOP (dbuf4) — nested groups x chunks (kept nested; flattening it + // with a conditional coopmat epilog crashes the Xclipse PAL compiler at + // large spec-resolved trip counts). One barrier per chunk. Chunk + // iteration (global index `chunk`): + // 1. barrier — A/B slice (chunk%2) fully written; on the first chunk + // of group g, wsc slice (g%2) is too. + // 2. prefetch — chunk+1 (A tiles, B blocks) into temp; when chunk+1 + // starts a new group, also its wsc element. Skipped + // entirely on the final chunk. + // 3. int8 MMA — on slice (chunk%2) into accum_int32. + // 4. store — temp -> A/B slice ((chunk+1)%2), unpacking the weight; + // on a group boundary, wsc -> slice ((g+1)%2). + // The group epilog runs unconditionally at the tail of each group. + // ========================================================= + uint chunk = 0; + for (uint group_i = 0; group_i < num_groups; ++group_i) { + for (uint inner = 0; inner < CHUNKS_PER_GROUP; ++inner, ++chunk) { + const bool has_next = chunk + 1u < num_chunks; + const bool group_crossing = has_next && (inner + 1u == CHUNKS_PER_GROUP); + const uint cur_a = (chunk % 2u) * ASH_SLICE_U32; + const uint cur_b = (chunk % 2u) * BSH_SLICE_U32; + const uint nxt_a = ((chunk + 1u) % 2u) * ASH_SLICE_U32; + const uint nxt_b = ((chunk + 1u) % 2u) * BSH_SLICE_U32; + + // coopmat-lds-fence: barrier() alone does NOT order shared stores against a + // subsequent coopMatLoad on the M51 Xclipse/AMD-PAL driver -- symptom is one + // stale MMA_M-row band of A, all columns, ~2.5% of runs, no crash (observed + // 2026-09-02 in sdpa_compute_out_coopmat.glsl). Measured cost of the fence: + // none (see this change's results). See memory + // `coopmat-lds-needs-explicit-memorybarriershared`. + memoryBarrierShared(); + barrier(); + + // --- 2. prefetch chunk+1 -> temp --- + if (has_next) { + const uint chunkK_nxt = (chunk + 1u) * WG_TILE_K; + // A staging (dbuf4tr's technique): coopMatLoad straight from global. + [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { + const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; + if (t < NUM_A_TILES) { + const uint tm = t / A_TILES_K; + const uint tk = t % A_TILES_K; + coopMatLoad( + temp_A[s], t_packed_int8_input, + (tile_m_start + tm * MMA_M) * a_row_stride_i8 + chunkK_nxt + + tk * MMA_K, + a_row_stride_i8, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#ifdef WEIGHT_INT4 + if (b_owns) { + [[unroll]] for (uint kk = 0; kk < (MMA_K >> 2u); ++kk) { + const uint k4_blk = (chunkK_nxt >> 2u) + b_k4base + kk; +#ifdef WEIGHT_BUFFER + temp_B[kk] = t_packed_weight[(b_n8blk1 * nblocks_x_A) + k4_blk]; +#else + temp_B[kk] = texelFetch(t_packed_weight, ivec2(k4_blk, b_n8blk1), 0); +#endif + } + } + if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { + const uint n_idx = tile_n_start + gl_LocalInvocationID.x; + f16vec4 sv = t_weight_scales[(group_i + 1u) * N4 + (n_idx >> 2u)]; + temp_wsc = float(sv[n_idx & 3u]); + } +#else + if (b_active) { + const uint block_y_w = (chunkK_nxt >> 2u) + b_k4_in_chunk; + const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; +#ifdef WEIGHT_BUFFER + temp_B = t_packed_weight[(block_y_w * N4) + block_x_w]; +#else + temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, block_y_w), 0); +#endif + } +#endif + } + + // --- 3. int8 MMA on the cur slice --- + [[unroll]] for (uint k = 0; k < NUM_K_SLABS; ++k) { + const uint slab_a_base_u32 = cur_a + k * A_SLAB_U32; + const uint slab_b_base_u32 = cur_b + k * B_SLAB_U32; + + coopmat matA[MMAS_PER_SG_M]; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + const uint row_a = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + coopMatLoad( + matA[i], Ash_int8, + slab_a_base_u32 + row_a * A_STRIDE_U32, + A_STRIDE_U32, + gl_CooperativeMatrixLayoutRowMajor); + } + + coopmat matB; + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint col_b = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopMatLoad( + matB, Bsh_int8, + slab_b_base_u32 + col_b * B_STRIDE_U32, + B_STRIDE_U32, + gl_CooperativeMatrixLayoutColumnMajor); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + accum_int32[i][j] = coopMatMulAdd(matA[i], matB, accum_int32[i][j]); + } + } + } + + // --- 4. store temp (chunk+1) -> nxt slice --- + if (has_next) { + // A staging (dbuf4tr's technique): coopMatStore into the nxt slice. + [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { + const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; + if (t < NUM_A_TILES) { + const uint tm = t / A_TILES_K; + const uint tk = t % A_TILES_K; + coopMatStore( + temp_A[s], Ash_int8, + nxt_a + tk * A_SLAB_U32 + (tm * MMA_M) * A_STRIDE_U32, + A_STRIDE_U32, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#ifdef WEIGHT_INT4 + if (b_owns) { + // 4 CONSECUTIVE uints from one thread -> merge-eligible into a wider + // store, versus 4 scalar stores from 4 different threads before. + [[unroll]] for (uint q = 0; q < (MMA_K >> 2u); ++q) { + const uint kk = (MMA_K >> 2u) - 1u - q; // DIAGNOSTIC: reverse order + Bsh_int8[nxt_b + b_lds_base + kk] = + widen_nibbles(uint(temp_B[kk][b_comp1]), b_par1); + } + } + if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { + const uint wbase_nxt = ((group_i + 1u) % 2u) * WG_TILE_N; + wsc_sh[wbase_nxt + gl_LocalInvocationID.x] = temp_wsc; + } +#else + if (b_active) { + const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); + const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); + const uint n_col_base = b_n_uint_col * 4u; + [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { + Bsh_int8[nxt_b + slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab] = + uint(temp_B[n_in_blk]); + } + } +#endif + } + } // chunks + + // --- Group epilog: scale-only accumulate, reset accum --- + // dbuf4zpg's, unchanged. Just result += float(acc) * wsc. The + // zero-point subtract and the ifs multiply are hoisted out of the group + // loop (applied once below). + { + const uint wbase = (group_i % 2u) * WG_TILE_N; + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + + coopmat wsc_bcast; + coopMatLoad( + wsc_bcast, wsc_sh, + wbase + local_n_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutRowMajor); + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + result[i][j] += + coopmat( + accum_int32[i][j]) * wsc_bcast; + accum_int32[i][j] = coopmat(0); + } + } + } + } // groups + + // --- Hoisted correction, applied ONCE: -------------------------------- + // result = ifs * ( result - izp * SUM_g wsc_g*wsum_g ) + // dbuf4zpg's, unchanged. izp/ifs are loaded here rather than before the + // group loop so they are not live across it. + { + coopmat + izpf_bcast[MMAS_PER_SG_M]; + coopmat + ifs_bcast[MMAS_PER_SG_M]; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + const uint local_m_base = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + coopmat izp_i; + coopMatLoad( + izp_i, izp_sh, local_m_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutColumnMajor); + izpf_bcast[i] = + coopmat(izp_i); + coopMatLoad( + ifs_bcast[i], ifs_sh, local_m_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutColumnMajor); + } + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopmat wcorr_bcast; + coopMatLoad( + wcorr_bcast, wcorr_sh, local_n_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutRowMajor); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + result[i][j] = ifs_bcast[i] * (result[i][j] - izpf_bcast[i] * wcorr_bcast); + } + } + } + + // --- Bias (optional) --- +#ifdef HAS_BIAS + if (apply_bias > 0) { + for (uint t = gl_LocalInvocationID.x; t < WG_TILE_N; t += WG_SIZE) { + bias_sh[t] = float(t_bias[tile_n_start + t]); + } + memoryBarrierShared(); + barrier(); + } +#endif + + // --- Store result tile --- + // N for the store address math MUST come from the spec constant, not the + // sizes UBO (see out_N_arg above). +#ifdef IO_TEXTURE + // Epilogue iteration i drains accumulator row-block i from EVERY subgroup + // into Csh_out at once, so the SG_GRID_Y bands it holds are disjoint global + // row ranges; the whole workgroup then imageStores them. lr / MMA_M is the + // writing subgroup's warpInTile.y, so the global row reproduces the buffer + // path's gi exactly. + // + // PORTABILITY NOTE: the barrier() in the loop body keeps this loop rolled + // despite [[unroll]], so result[i][j] IS dynamically indexed. Coopmat arrays + // are opaque per-lane storage and dynamic indexing is exactly the construct + // the Xclipse/AMD-PAL compiler has broken before -- check this first if the + // texture variants miscompile on M51. + const uint CSH_TEXELS_PER_ROW = WG_TILE_N / 4u; + const uint CSH_TEXELS = CSH_ROWS * CSH_TEXELS_PER_ROW; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + // Guards Csh_out against the previous iteration's readers. Inert on i == 0 + // but must stay unconditional to remain workgroup-uniform. + // coopmat-lds-fence: barrier() alone does NOT order shared stores against a + // subsequent coopMatLoad on the M51 Xclipse/AMD-PAL driver -- symptom is one + // stale MMA_M-row band of A, all columns, ~2.5% of runs, no crash (observed + // 2026-09-02 in sdpa_compute_out_coopmat.glsl). Measured cost of the fence: + // none (see this change's results). See memory + // `coopmat-lds-needs-explicit-memorybarriershared`. + memoryBarrierShared(); + barrier(); + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { +#ifdef HAS_BIAS + if (apply_bias > 0) { + const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopmat bias_tile; + coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); + result[i][j] += bias_tile; + } +#endif + coopmat out_tile = + coopmat(result[i][j]); + coopMatStore( + out_tile, Csh_out, + warpInTile.y * MMA_M * WG_TILE_N + + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j), + WG_TILE_N, + gl_CooperativeMatrixLayoutRowMajor); + } + memoryBarrierShared(); + barrier(); + + for (uint t = gl_LocalInvocationID.x; t < CSH_TEXELS; t += WG_SIZE) { + const uint lr = t / CSH_TEXELS_PER_ROW; + const uint lc4 = t % CSH_TEXELS_PER_ROW; + const uint m = + tile_m_start + (lr / MMA_M) * SG_TILE_M + i * MMA_M + (lr % MMA_M); + const uint base = lr * WG_TILE_N + lc4 * 4u; + imageStore( + t_output, + ivec3(tile_n_start / 4u + lc4, m, 0), + vec4( + float(Csh_out[base]), + float(Csh_out[base + 1u]), + float(Csh_out[base + 2u]), + float(Csh_out[base + 3u]))); + } + } +#else + const uint N_out = uint(out_N_arg); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint gi = tile_m_start + MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + const uint gj = tile_n_start + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + +#ifdef HAS_BIAS + if (apply_bias > 0) { + const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopmat bias_tile; + coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); + result[i][j] += bias_tile; + } +#endif + + coopmat out_tile = + coopmat(result[i][j]); + coopMatStore( + out_tile, t_output, + gi * N_out + gj, N_out, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#endif // IO_TEXTURE +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbwr.yaml b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbwr.yaml new file mode 100644 index 00000000000..e7370ebafd2 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbwr.yaml @@ -0,0 +1,68 @@ +# !!! KNOWN-INCORRECT -- DO NOT PROMOTE, DO NOT TIME !!! +# Fails the correctness gate deterministically (3/3) on the num_groups==2 +# shapes only (K=256). Cause isolated 2026-09-08: the 4-wide-per-thread +# ownership with half the workgroup idle. Confirmed by dbuf4zpgbw2, which +# uses the SAME consecutive-write idea at 2-wide with all 256 threads +# active and passes 14/14 x3. Store ORDER is not the trigger (dbuf4zpgbwr, +# reversed, fails identically) and neither is the uvec4 retype (dbuf4zpgbw, +# no retype, fails identically). +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# "zpg" + "-tr" combination: tsweep_dbuf4zpg with its per-thread scalar +# A-staging replaced by tsweep_dbuf4tr's coopMat-mediated A staging. B +# staging, zp-hoist, byte-parallel nibble widening are all dbuf4zpg's, +# unchanged -- only A-staging differs. Requires t_packed_int8_input in the +# ROW-MAJOR kPackedInt8_4W layout (same requirement as tsweep_dbuf4tr). +# +# PROMOTED 2026-09-01 as the shipped default -- see QuantizedLinear.cpp's +# dq8ca_coopmat_variant() for the full validation record (10/10 buffer, 6/6 +# texture3d correctness across 1B/3B/8B; +4.2% real e2e prefill; 46.50% -> +# 49.94% efficiency on 8B). Also selectable explicitly via +# ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4zpgbwr_txkgs. +# See openspec/changes/archive/2026-08-31-dq8ca-tr-staged-a-on-zpg. +# +# Single seed tile: the prior dbuf4zpg default's own tile (t128x64k32g42s32). +# A re-sweep against this shader's own (lower) register-pressure profile was +# run as a follow-up (coopmat-tr-tilesweep-4w-port) and found no better tile +# -- this remains the best known geometry. + +linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbwr: + parameter_names_with_default_values: + PRECISION: highp + HAS_BIAS: false + IO_STORAGE: buffer + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + MMA_M: 16 + MMA_N: 16 + MMA_K: 16 + WG_TILE_M: 64 + WG_TILE_N: 32 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + shader_variants: + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbwr_t128x64k32g42s32_buffer_texture2d_half + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbwr_t128x64k32g42s32_texture3d_texture2d_half + IO_STORAGE: texture3d + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr3.glsl b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr3.glsl new file mode 100644 index 00000000000..24d42f12eff --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr3.glsl @@ -0,0 +1,781 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +/* + * "zpg" + "-tr" combination: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpg.glsl + * with its per-thread scalar A-staging replaced by + * linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr.glsl's coopMat-mediated A + * staging (coopMatLoad(global) -> coopmat<> -> coopMatStore(LDS)), PLUS tr3's + * shared-A layout: scalar int8_t storage, row-major over the FULL chunk, with + * A_ROW_PAD_I8 bytes of per-row padding. B staging, the quantization math and + * every correctness safeguard are dbuf4zpgtr's, byte-for-byte unchanged. This is + * an ADDITIVE combination, not a redesign: every non-A-staging block below + * (B staging: coalesced write, no skew; zp-hoist: izp/ifs applied once after + * the group loop via wcorr_sh; byte-parallel nibble widening; static + * A_ALWAYS_ACTIVE branch elision -- N/A here, see below; group epilog; + * bias/store epilogue) is byte-identical to dbuf4zpg's. Only the A-staging + * block (prologue load+store, main-loop prefetch+store) is dbuf4tr's, + * verbatim. + * + * dbuf4zpg's per-thread A staging used an `a_active` guard (statically always + * true when A_ACTIVE_THREADS == WG_SIZE, via the A_MAP_FULL-gated + * A_ALWAYS_ACTIVE macro). dbuf4tr's per-SUBGROUP tile map has no equivalent + * concept -- every subgroup participates via a `t < NUM_A_TILES` guard that + * depends only on gl_SubgroupID, not gl_LocalInvocationID.x -- so + * A_MAP_FULL/A_ALWAYS_ACTIVE is dropped entirely in this file; it would be + * dead code for the new A-staging block. + * + * Rationale for combining this way (not the reverse) and why this is worth + * building at all: see this change's design.md D0-D3. In short -- the only + * existing measurement of dbuf4tr's A-staging technique (28.72-30.51%, + * dq8ca-arch-redesign) was taken against dbuf4tr's own pre-zpg baseline + * (old B skew, no byte-parallel widening, no branch elision) -- a materially + * weaker shader than the 46.49-46.50% dbuf4zpg this file now combines it + * with. This file exists to answer whether that combination performs + * differently now that register pressure is already reduced. + * + * A staging (the actual delta from dbuf4zpg): + * dbuf4zpg: per-thread (m4, k4) ivec4 fetch, hoisted a_lds_off0/a_glb_row; + * only A_ACTIVE_THREADS invocations participate, each scattering + * 4 rows into shared A with 4 scalar stores. + * this file: per-SUBGROUP MMA_M x MMA_K tile fetch via coopMatLoad straight + * from a ROW-MAJOR (kPackedInt8_4W) int8 activation buffer, then + * coopMatStore into shared A -- dbuf4tr's mapping, + * unmodified (not re-derived; see design.md D3). + * + * t_packed_int8_input is therefore bound the same way dbuf4tr binds it: a + * SCALAR int8_t array in the kPackedInt8_4W layout (plain row-major int8, + * row stride K), produced by quantize_and_pack_4w_with_group_sums.glsl. + * QuantizedLinear.cpp's dq8ca_variant_wants_rowmajor_a() must recognize this + * file's variant token (tsweep_dbuf4zpgtr_t...) the same way it already + * recognizes tsweep_dbuf4tr_t/trm_t/trd_t, so graph-build time (packer + * selection) and dispatch time (kernel selection) cannot disagree. + * + * B CANNOT be coopmat-staged (int4 nibble unpack; a coopmat's per-lane layout + * is opaque to hand-assembly from unpacked registers) -- unchanged from both + * parent files. B staging below is dbuf4zpg's byte-parallel, coalesced, + * no-skew version, untouched. + * + * The loop structure is dbuf4's (both parents share it), unchanged: + * prologue: prefetch chunk 0 -> temp, store to slice 0 (no barrier) + * per iter: barrier -> prefetch(next) -> MMA(cur) -> store(next) + * kept nested (groups x chunks) with an unconditional group epilog -- + * flattening it crashes the Xclipse PAL compiler at large spec-resolved trip + * counts (see dbuf2's own header). + * + * Selected via + * ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4zpgtr_txkgs<32|64> + * (QuantizedLinear.cpp), additive to the tsweep_dbuf4zpg_t..., tsweep_dbuf4tr_t... + * and tsweep_t... namespaces. NOT the default -- unvalidated until it passes + * repeated test_llama_microbench --correctness-only runs (see + * dq8ca_coopmat_variant()'s comment on why a single pass is not proof). + * + * Performs: out[M,N] = dequant(int8_act) * dequant(int_w) (+ bias) + * via coopmat x coopmat -> coopmat on the matrix unit. + * + * Hard preconditions (dbuf4zpg's, plus dbuf4tr's row-major/alignment ones): + * M % WG_TILE_M == 0, N % WG_TILE_N == 0, K % WG_TILE_K == 0, + * group_size % WG_TILE_K == 0, K % 4 == 0, + * WG_TILE_M % MMA_M == 0, WG_TILE_K % MMA_K == 0, + * t_packed_int8_input in kPackedInt8_4W (row-major) layout, + * device exposes coopmatx-> at 16x16x16. + */ + +#version 450 core + +#extension GL_KHR_cooperative_matrix : require +#extension GL_KHR_memory_scope_semantics : require +#extension GL_KHR_shader_subgroup_basic : enable +#extension GL_EXT_shader_explicit_arithmetic_types : require +#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require +// 8-bit SSBO access: A is bound as a scalar int8_t array so that the +// coopMatLoad below has a MATCHING component type (see dbuf4tr's header for +// why the type must match on this driver). +#extension GL_EXT_shader_8bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_EXT_control_flow_attributes : enable + +#define PRECISION ${PRECISION} + +$if WEIGHT_NBITS == 4: + #define WEIGHT_INT4 + +$if HAS_BIAS: + #define HAS_BIAS + +$if WEIGHT_STORAGE == "buffer": + #define WEIGHT_BUFFER + +$if IO_STORAGE == "texture3d": + #define IO_TEXTURE + +layout(std430) buffer; + +#include "common.glslh" + +// Bindings — match add_linear_dqa_qw_node arg order: +// output(0), fp_input(1), packed_int8_input(2), int_input_sums(3 - unused), +// input_scales(4), input_zps(5), packed_weight(6), weight_sums(7), +// weight_scales(8), bias(9). +${layout_declare_tensor(B, "w", "t_output", "half", IO_STORAGE, is_scalar_array=True)} +// t_input is unread here -- the activations arrive already quantized in +// t_packed_int8_input -- but stays declared so the binding layout matches the +// dispatch site. It tracks IO_STORAGE so the two IO tensors stay consistent. +${layout_declare_tensor(B, "r", "t_input", "half", IO_STORAGE, is_scalar_array=False)} +// ROW-MAJOR (kPackedInt8_4W) packed activations, bound as a scalar int8_t +// array (row stride = K int8) -- dbuf4tr's binding, unchanged. The stock +// 4h4w layout dbuf4zpg uses is NOT row-major (component index selects a row, +// non-affine), so it cannot be addressed by any coopMatLoad. +${layout_declare_tensor(B, "r", "t_packed_int8_input", "int8", "buffer", is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_int8_input_sums", "int", "buffer", is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_int8_input_scales", "half", "texture3d")} +${layout_declare_tensor(B, "r", "t_int8_input_zps", "int8", "texture3d")} +${layout_declare_tensor(B, "r", "t_packed_weight", "int", WEIGHT_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_weight_sums", "int", "buffer", is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_weight_scales", "half", "buffer", is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_bias", "half", "buffer", is_scalar_array=True)} + +${layout_declare_ubo(B, "ivec4", "output_sizes")} +${layout_declare_ubo(B, "ivec4", "input_sizes")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +${layout_declare_spec_const(C, "int", "apply_bias", "0")} +// INT4 only; inert (0) for INT8 so the dispatcher's spec list lines up. +${layout_declare_spec_const(C, "int", "K4_per_group", "0")} +${layout_declare_spec_const(C, "int", "num_groups_arg", "0")} +${layout_declare_spec_const(C, "int", "out_N_arg", "0")} + +// Tile geometry +const uint MMA_M = ${MMA_M}; +const uint MMA_N = ${MMA_N}; +const uint MMA_K = ${MMA_K}; + +const uint WG_TILE_M = ${WG_TILE_M}; +const uint WG_TILE_N = ${WG_TILE_N}; +const uint WG_TILE_K = ${WG_TILE_K}; + +const uint SG_GRID_X = ${SG_GRID_X}; +const uint SG_GRID_Y = ${SG_GRID_Y}; +const uint SUBGROUP_SIZE = ${SUBGROUP_SIZE}; +const uint NUM_SUBGROUPS = SG_GRID_X * SG_GRID_Y; +const uint WG_SIZE = NUM_SUBGROUPS * SUBGROUP_SIZE; + +const uint SG_TILE_M = WG_TILE_M / SG_GRID_Y; +const uint SG_TILE_N = WG_TILE_N / SG_GRID_X; +const uint MMAS_PER_SG_M = SG_TILE_M / MMA_M; +const uint MMAS_PER_SG_N = SG_TILE_N / MMA_N; + +const uint B_USEFUL_U32 = MMA_K / 4u; +// No skew + coalesced write -- dbuf4zpg's B fix, unchanged (this file does +// not touch B staging at all). +const uint B_STRIDE_U32 = B_USEFUL_U32; +const uint B_SLAB_U32 = WG_TILE_N * B_STRIDE_U32; +const uint NUM_K_SLABS = WG_TILE_K / MMA_K; + +// ===== tr3-style A shared layout: THIS FILE'S ONLY STRUCTURAL CHANGE ===== +// dbuf4zpgtr stages A slab-major -- one MMA_K-wide slab per K step, row stride +// MMA_K/4 dwords, no padding, element type `uint` (4 int8 packed per element). +// tr3 instead keeps ONE row-major row per A row spanning the WHOLE chunk +// (WG_TILE_K wide), pads that row, and declares the array with the coopmat's +// own scalar element type so coopMatLoad/Store address it in elements. +// +// A_ROW_PAD_I8 pads the FULL-CHUNK row, NOT each MMA_K slab row -- the two have +// very different storage costs. Full-chunk padding costs WG_TILE_M*A_ROW_PAD_I8 +// bytes per slice; padding every slab row would cost NUM_K_SLABS times that. +// tr3's own value is 16 bytes (== ELEMENTS_PER_VEC4 for an 8-bit type). +const uint A_ROW_PAD_I8 = ${A_ROW_PAD_I8}u; +const uint STRIDE_A_I8 = WG_TILE_K + A_ROW_PAD_I8; +const uint ASH_SLICE_I8 = WG_TILE_M * STRIDE_A_I8; + +const uint BSH_SLICE_U32 = NUM_K_SLABS * B_SLAB_U32; + +// Double-buffered MMA operand staging. A is scalar int8_t (tr3-style); B keeps +// dbuf4zpg's packed-uint column-major slab layout untouched at this stage. +shared int8_t Ash_i8[2u * ASH_SLICE_I8]; +shared uint Bsh_int8[2u * BSH_SLICE_U32]; + +// Per-WG-tile-row activation params (loaded ONCE at WG start; constant +// across groups). +shared int izp_sh[WG_TILE_M]; // int32 (cast from int8 source) for broadcast +shared float ifs_sh[WG_TILE_M]; // float32 (cast from fp16 source) for broadcast + +// Per-(group, output-channel) weight params, ping-ponged by group parity. +// (For per-channel INT8 only slice 0 is ever used.) +shared float wsc_sh[2u * WG_TILE_N]; +// SUM_g wsc[g][n]*wsum[g][n] per output channel -- weight-side only, so it is +// accumulated once in the prologue. dbuf4zpg's zp-hoist, unchanged. +shared float wcorr_sh[WG_TILE_N]; + +#ifdef HAS_BIAS +shared float bias_sh[WG_TILE_N]; +#endif + +#ifdef IO_TEXTURE +// Result staging for the imageStore epilogue, mirroring the fp16 kernel: +// SG_GRID_Y bands of MMA_M rows, each WG_TILE_N wide, row-major. A full +// WG_TILE_M x WG_TILE_N buffer would cost SG_GRID_Y/MMAS_PER_SG_M x more LDS +// and wreck occupancy. float16_t-typed because coopMatStore needs it. +const uint CSH_ROWS = SG_GRID_Y * MMA_M; +shared float16_t Csh_out[CSH_ROWS * WG_TILE_N]; +#endif + +// Running fp32 accumulator (across all groups). +coopmat + result[MMAS_PER_SG_M][MMAS_PER_SG_N]; + +// Per-group int32 MMA accumulator. +coopmat + accum_int32[MMAS_PER_SG_M][MMAS_PER_SG_N]; + + +// Byte-parallel int4 -> int8 widening. dbuf4zpg's, unchanged (B-side only). +// +// The four nibbles this shader needs from one packed uint are ALREADY one per +// byte (bits 3:0 of each byte for parity 0, bits 7:4 for parity 1), so all four +// can be widened at once instead of with a per-nibble +// shift/mask/bias-subtract/mask chain. +// +// For v in [0,15] the biased value is v-8. `v ^ 8` is exactly the 4-bit +// two's-complement pattern of v-8, because -8 == +8 (mod 16): +// v=0 -> 0x8 -> -8 v=7 -> 0xF -> -1 +// v=8 -> 0x0 -> 0 v=15 -> 0x7 -> +7 +// so the only remaining work is sign-extending bit 3 into bits 7:4 per byte. +// `sgn * 0x1E` does that with no cross-byte carry: 0x08 * 0x1E == 0xF0 exactly, +// and sgn is at most 0x08080808 so the product is at most 0xF0F0F0F0. +// +// A naive `nib - 0x08080808` would NOT work -- it borrows across byte lanes +// whenever a nibble is < 8. Shifts must be on uint, not int, so the >> is +// logical rather than arithmetic. +// +// ~5 ops per 4 weights vs ~22 for the per-nibble chain; bit-identical output. +uint widen_nibbles(const uint w, const uint parity) { + const uint nib = (parity == 0u) ? (w & 0x0F0F0F0Fu) : ((w >> 4u) & 0x0F0F0F0Fu); + const uint p = nib ^ 0x08080808u; + const uint sgn = p & 0x08080808u; + return p | (sgn * 0x1Eu); +} + +void main() { + const uvec2 tileID = uvec2(gl_WorkGroupID.xy); + const uvec2 warpInTile = uvec2( + gl_SubgroupID % SG_GRID_X, + gl_SubgroupID / SG_GRID_X); + + const uint K = uint(input_sizes.x); + const uint N = uint(output_sizes.x); + const uint N4 = (N + 3u) / 4u; + const uint nblocks_x_A = (K + 3u) >> 2u; + // A row stride in INT8 elements (dbuf4tr's binding is row-major int8, not + // the 4h4w ivec4 block layout dbuf4zpg used -- so A addressing below is in + // int8 elements, not int, and derived from nblocks_x_A so it matches the + // packer's `m_row * K4 + k4` addressing exactly (K % 4 == 0 makes them + // equal to K directly). + const uint a_row_stride_i8 = nblocks_x_A * 4u; + +#ifdef WEIGHT_INT4 + const uint num_groups = uint(num_groups_arg); + const uint CHUNKS_PER_GROUP = uint(K4_per_group) * 4u / WG_TILE_K; +#else + // Per-channel: a single quant "group" spanning all of K. The nested + // groups x chunks loop below collapses to a flat chunk loop, the wsum/wsc + // ping-pong never crosses a boundary, and the epilog runs exactly once. + const uint num_groups = 1u; + const uint CHUNKS_PER_GROUP = uint(num_groups_arg); +#endif + const uint num_chunks = num_groups * CHUNKS_PER_GROUP; + + const uint tile_m_start = WG_TILE_M * tileID.y; + const uint tile_n_start = WG_TILE_N * tileID.x; + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + result[i][j] = coopmat(0.0); + accum_int32[i][j] = coopmat(0); + } + } + + const uint K_BLOCKS_PER_CHUNK = WG_TILE_K >> 2u; + + // --- A staging tile map (dbuf4tr's, unmodified): one MMA_M x MMA_K coopmat + // tile per subgroup per slot, dealt round-robin across the + // NUM_SUBGROUPS subgroups so every subgroup participates. Replaces + // dbuf4zpg's per-thread (m4, k4) map / a_active guard entirely -- see + // design.md D3 for why this is reused as-is, not re-derived. + const uint A_TILES_M = WG_TILE_M / MMA_M; + const uint A_TILES_K = WG_TILE_K / MMA_K; // == NUM_K_SLABS + const uint NUM_A_TILES = A_TILES_M * A_TILES_K; + const uint A_TILES_PER_SG = (NUM_A_TILES + NUM_SUBGROUPS - 1u) / NUM_SUBGROUPS; + +#ifdef WEIGHT_INT4 + // --- B staging thread map: (block, col) slots; each slot extracts one + // ColumnMajor LDS uint (4 K-contiguous sign-extended int8) --- + const uint B_TOTAL_SLOTS = K_BLOCKS_PER_CHUNK * WG_TILE_N; + const uint B_SLOTS_PER_THREAD = B_TOTAL_SLOTS / WG_SIZE; + const uint N8_PER_TILE = WG_TILE_N >> 3u; +#else + // --- B staging thread map: one (k4, n4) ivec4 block per active thread --- + const uint B_FETCH_SLOTS = K_BLOCKS_PER_CHUNK * (WG_TILE_N >> 2u); + const uint N4_PER_TILE = WG_TILE_N >> 2u; + const uint b_k4_in_chunk = gl_LocalInvocationID.x / N4_PER_TILE; + const uint b_n_uint_col = gl_LocalInvocationID.x % N4_PER_TILE; + const bool b_active = gl_LocalInvocationID.x < B_FETCH_SLOTS; +#endif + + // ===== INTERVENTION F: hoist loop-invariant B staging index math ===== + // dbuf4zpg's, unchanged -- B staging is untouched by this file's A-staging + // swap. See dbuf4zpg's header for the full rationale (ablation-attributed + // -16.8% block, mostly index arithmetic not memory traffic). +#ifdef WEIGHT_INT4 + uint b_lds_off[B_SLOTS_PER_THREAD]; // LDS store offset within a slice + uint b_comp[B_SLOTS_PER_THREAD]; // which ivec4 component feeds this slot + uint b_par[B_SLOTS_PER_THREAD]; // nibble parity for this slot + uint b_n8blk[B_SLOTS_PER_THREAD]; // global texel column (N/8 blocks) + uint b_k4off[B_SLOTS_PER_THREAD]; // k4 offset of this slot within a chunk + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { + const uint a = gl_LocalInvocationID.x + si * WG_SIZE; + const uint slab_idx = a / B_SLAB_U32; + const uint local_a = a % B_SLAB_U32; + const uint n_col = local_a / B_STRIDE_U32; + const uint k4_in_slab = local_a % B_STRIDE_U32; + const uint k4_in_chunk = slab_idx * (MMA_K >> 2u) + k4_in_slab; + const uint n8_in_tile = n_col >> 3u; + const uint rem = n_col & 7u; + b_lds_off[si] = a; + b_comp[si] = rem & 3u; + b_par[si] = rem >> 2u; + b_n8blk[si] = (tile_n_start >> 3u) + n8_in_tile; + b_k4off[si] = k4_in_chunk; + } +#endif + + // Prefetch temp registers. temp_A is a coopmat array (dbuf4tr's A-staging + // technique); indices into it are [[unroll]]-resolved compile-time + // constants, never dynamic -- dynamic indexing of a coopmat array is + // exactly the construct the Xclipse/AMD-PAL compiler has miscompiled + // before. + coopmat + temp_A[A_TILES_PER_SG]; +#ifdef WEIGHT_INT4 + ivec4 temp_B[B_SLOTS_PER_THREAD]; + float temp_wsc; +#else + ivec4 temp_B; +#endif + + // ========================================================= + // PROLOGUE + // ========================================================= + if (gl_LocalInvocationID.x < (WG_TILE_M >> 2u)) { + const uint m4 = (tile_m_start >> 2u) + gl_LocalInvocationID.x; + const vec4 sc = vec4(texelFetch(t_int8_input_scales, ivec3(m4, 0, 0), 0)); + const ivec4 zp = texelFetch(t_int8_input_zps, ivec3(m4, 0, 0), 0); + const uint base = gl_LocalInvocationID.x * 4u; + ifs_sh[base + 0u] = sc.x; ifs_sh[base + 1u] = sc.y; + ifs_sh[base + 2u] = sc.z; ifs_sh[base + 3u] = sc.w; + izp_sh[base + 0u] = zp.x; izp_sh[base + 1u] = zp.y; + izp_sh[base + 2u] = zp.z; izp_sh[base + 3u] = zp.w; + } + // Group 0 weight scales -> slice 0, and the hoisted weight-side correction + // SUM_g wsc[g][n]*wsum[g][n] accumulated across ALL groups. dbuf4zpg's + // zp-hoist, unchanged. + if (gl_LocalInvocationID.x < WG_TILE_N) { + const uint n_idx = tile_n_start + gl_LocalInvocationID.x; + f16vec4 sv0 = t_weight_scales[n_idx >> 2u]; + wsc_sh[gl_LocalInvocationID.x] = float(sv0[n_idx & 3u]); + + float corr = 0.0; + for (uint g = 0; g < num_groups; ++g) { + f16vec4 sv = t_weight_scales[g * N4 + (n_idx >> 2u)]; + corr += float(sv[n_idx & 3u]) * float(t_weight_sums[g * N + n_idx]); + } + wcorr_sh[gl_LocalInvocationID.x] = corr; + } + memoryBarrierShared(); + barrier(); + + // NOTE: dbuf4zpg builds izp_bcast/ifs_bcast AFTER the group loop, not here + // -- that is the register-pressure saving zp-hoist buys. Unchanged. + + // dbuf4: prefetch chunk 0 into temp registers, THEN store to slice 0 (no + // barrier here -- the main loop's first iteration barriers before + // reading slice 0). + // + // A staging (dbuf4tr's technique): per-subgroup coopMatLoad straight from + // the row-major global buffer. + [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { + const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; + if (t < NUM_A_TILES) { + const uint tm = t / A_TILES_K; + const uint tk = t % A_TILES_K; + coopMatLoad( + temp_A[s], t_packed_int8_input, + (tile_m_start + tm * MMA_M) * a_row_stride_i8 + tk * MMA_K, + a_row_stride_i8, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#ifdef WEIGHT_INT4 + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { +#ifdef WEIGHT_BUFFER + temp_B[si] = t_packed_weight[(b_n8blk[si] * nblocks_x_A) + b_k4off[si]]; +#else + temp_B[si] = texelFetch(t_packed_weight, ivec2(b_k4off[si], b_n8blk[si]), 0); +#endif + } +#else + if (b_active) { + const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; +#ifdef WEIGHT_BUFFER + temp_B = t_packed_weight[(b_k4_in_chunk * N4) + block_x_w]; +#else + temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, b_k4_in_chunk), 0); +#endif + } +#endif + { + // store chunk 0 -> slice 0 + // A staging (dbuf4tr's technique): coopMatStore into the same Ash_int8 + // slot layout dbuf4zpg's scalar scatter used to write. + [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { + const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; + if (t < NUM_A_TILES) { + const uint tm = t / A_TILES_K; + const uint tk = t % A_TILES_K; + coopMatStore( + temp_A[s], Ash_i8, + (tm * MMA_M) * STRIDE_A_I8 + tk * MMA_K, + STRIDE_A_I8, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#ifdef WEIGHT_INT4 + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { + Bsh_int8[b_lds_off[si]] = + widen_nibbles(uint(temp_B[si][b_comp[si]]), b_par[si]); + } +#else + if (b_active) { + const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); + const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); + const uint n_col_base = b_n_uint_col * 4u; + [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { + Bsh_int8[slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab] = + uint(temp_B[n_in_blk]); + } + } +#endif + } + + // ========================================================= + // MAIN LOOP (dbuf4) — nested groups x chunks (kept nested; flattening it + // with a conditional coopmat epilog crashes the Xclipse PAL compiler at + // large spec-resolved trip counts). One barrier per chunk. Chunk + // iteration (global index `chunk`): + // 1. barrier — A/B slice (chunk%2) fully written; on the first chunk + // of group g, wsc slice (g%2) is too. + // 2. prefetch — chunk+1 (A tiles, B blocks) into temp; when chunk+1 + // starts a new group, also its wsc element. Skipped + // entirely on the final chunk. + // 3. int8 MMA — on slice (chunk%2) into accum_int32. + // 4. store — temp -> A/B slice ((chunk+1)%2), unpacking the weight; + // on a group boundary, wsc -> slice ((g+1)%2). + // The group epilog runs unconditionally at the tail of each group. + // ========================================================= + uint chunk = 0; + for (uint group_i = 0; group_i < num_groups; ++group_i) { + for (uint inner = 0; inner < CHUNKS_PER_GROUP; ++inner, ++chunk) { + const bool has_next = chunk + 1u < num_chunks; + const bool group_crossing = has_next && (inner + 1u == CHUNKS_PER_GROUP); + const uint cur_a = (chunk % 2u) * ASH_SLICE_I8; + const uint cur_b = (chunk % 2u) * BSH_SLICE_U32; + const uint nxt_a = ((chunk + 1u) % 2u) * ASH_SLICE_I8; + const uint nxt_b = ((chunk + 1u) % 2u) * BSH_SLICE_U32; + + // coopmat-lds-fence: barrier() alone does NOT order shared stores against a + // subsequent coopMatLoad on the M51 Xclipse/AMD-PAL driver -- symptom is one + // stale MMA_M-row band of A, all columns, ~2.5% of runs, no crash (observed + // 2026-09-02 in sdpa_compute_out_coopmat.glsl). Measured cost of the fence: + // none (see this change's results). See memory + // `coopmat-lds-needs-explicit-memorybarriershared`. + memoryBarrierShared(); + barrier(); + + // --- 2. prefetch chunk+1 -> temp --- + if (has_next) { + const uint chunkK_nxt = (chunk + 1u) * WG_TILE_K; + // A staging (dbuf4tr's technique): coopMatLoad straight from global. + [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { + const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; + if (t < NUM_A_TILES) { + const uint tm = t / A_TILES_K; + const uint tk = t % A_TILES_K; + coopMatLoad( + temp_A[s], t_packed_int8_input, + (tile_m_start + tm * MMA_M) * a_row_stride_i8 + chunkK_nxt + + tk * MMA_K, + a_row_stride_i8, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#ifdef WEIGHT_INT4 + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { + const uint k4_blk = (chunkK_nxt >> 2u) + b_k4off[si]; +#ifdef WEIGHT_BUFFER + temp_B[si] = t_packed_weight[(b_n8blk[si] * nblocks_x_A) + k4_blk]; +#else + temp_B[si] = texelFetch(t_packed_weight, ivec2(k4_blk, b_n8blk[si]), 0); +#endif + } + if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { + const uint n_idx = tile_n_start + gl_LocalInvocationID.x; + f16vec4 sv = t_weight_scales[(group_i + 1u) * N4 + (n_idx >> 2u)]; + temp_wsc = float(sv[n_idx & 3u]); + } +#else + if (b_active) { + const uint block_y_w = (chunkK_nxt >> 2u) + b_k4_in_chunk; + const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; +#ifdef WEIGHT_BUFFER + temp_B = t_packed_weight[(block_y_w * N4) + block_x_w]; +#else + temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, block_y_w), 0); +#endif + } +#endif + } + + // --- 3. int8 MMA on the cur slice --- + [[unroll]] for (uint k = 0; k < NUM_K_SLABS; ++k) { + // row-major full-chunk A: the K step is a column offset within the + // row, not a separate slab base. + const uint k_col_a_i8 = k * MMA_K; + const uint slab_b_base_u32 = cur_b + k * B_SLAB_U32; + + coopmat matA[MMAS_PER_SG_M]; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + const uint row_a = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + coopMatLoad( + matA[i], Ash_i8, + cur_a + row_a * STRIDE_A_I8 + k_col_a_i8, + STRIDE_A_I8, + gl_CooperativeMatrixLayoutRowMajor); + } + + coopmat matB; + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint col_b = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopMatLoad( + matB, Bsh_int8, + slab_b_base_u32 + col_b * B_STRIDE_U32, + B_STRIDE_U32, + gl_CooperativeMatrixLayoutColumnMajor); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + accum_int32[i][j] = coopMatMulAdd(matA[i], matB, accum_int32[i][j]); + } + } + } + + // --- 4. store temp (chunk+1) -> nxt slice --- + if (has_next) { + // A staging (dbuf4tr's technique): coopMatStore into the nxt slice. + [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { + const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; + if (t < NUM_A_TILES) { + const uint tm = t / A_TILES_K; + const uint tk = t % A_TILES_K; + coopMatStore( + temp_A[s], Ash_i8, + nxt_a + (tm * MMA_M) * STRIDE_A_I8 + tk * MMA_K, + STRIDE_A_I8, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#ifdef WEIGHT_INT4 + [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { + Bsh_int8[nxt_b + b_lds_off[si]] = + widen_nibbles(uint(temp_B[si][b_comp[si]]), b_par[si]); + } + if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { + const uint wbase_nxt = ((group_i + 1u) % 2u) * WG_TILE_N; + wsc_sh[wbase_nxt + gl_LocalInvocationID.x] = temp_wsc; + } +#else + if (b_active) { + const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); + const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); + const uint n_col_base = b_n_uint_col * 4u; + [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { + Bsh_int8[nxt_b + slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab] = + uint(temp_B[n_in_blk]); + } + } +#endif + } + } // chunks + + // --- Group epilog: scale-only accumulate, reset accum --- + // dbuf4zpg's, unchanged. Just result += float(acc) * wsc. The + // zero-point subtract and the ifs multiply are hoisted out of the group + // loop (applied once below). + { + const uint wbase = (group_i % 2u) * WG_TILE_N; + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + + coopmat wsc_bcast; + coopMatLoad( + wsc_bcast, wsc_sh, + wbase + local_n_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutRowMajor); + + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + result[i][j] += + coopmat( + accum_int32[i][j]) * wsc_bcast; + accum_int32[i][j] = coopmat(0); + } + } + } + } // groups + + // --- Hoisted correction, applied ONCE: -------------------------------- + // result = ifs * ( result - izp * SUM_g wsc_g*wsum_g ) + // dbuf4zpg's, unchanged. izp/ifs are loaded here rather than before the + // group loop so they are not live across it. + { + coopmat + izpf_bcast[MMAS_PER_SG_M]; + coopmat + ifs_bcast[MMAS_PER_SG_M]; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + const uint local_m_base = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + coopmat izp_i; + coopMatLoad( + izp_i, izp_sh, local_m_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutColumnMajor); + izpf_bcast[i] = + coopmat(izp_i); + coopMatLoad( + ifs_bcast[i], ifs_sh, local_m_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutColumnMajor); + } + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopmat wcorr_bcast; + coopMatLoad( + wcorr_bcast, wcorr_sh, local_n_base, /*stride=*/0u, + gl_CooperativeMatrixLayoutRowMajor); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + result[i][j] = ifs_bcast[i] * (result[i][j] - izpf_bcast[i] * wcorr_bcast); + } + } + } + + // --- Bias (optional) --- +#ifdef HAS_BIAS + if (apply_bias > 0) { + for (uint t = gl_LocalInvocationID.x; t < WG_TILE_N; t += WG_SIZE) { + bias_sh[t] = float(t_bias[tile_n_start + t]); + } + memoryBarrierShared(); + barrier(); + } +#endif + + // --- Store result tile --- + // N for the store address math MUST come from the spec constant, not the + // sizes UBO (see out_N_arg above). +#ifdef IO_TEXTURE + // Epilogue iteration i drains accumulator row-block i from EVERY subgroup + // into Csh_out at once, so the SG_GRID_Y bands it holds are disjoint global + // row ranges; the whole workgroup then imageStores them. lr / MMA_M is the + // writing subgroup's warpInTile.y, so the global row reproduces the buffer + // path's gi exactly. + // + // PORTABILITY NOTE: the barrier() in the loop body keeps this loop rolled + // despite [[unroll]], so result[i][j] IS dynamically indexed. Coopmat arrays + // are opaque per-lane storage and dynamic indexing is exactly the construct + // the Xclipse/AMD-PAL compiler has broken before -- check this first if the + // texture variants miscompile on M51. + const uint CSH_TEXELS_PER_ROW = WG_TILE_N / 4u; + const uint CSH_TEXELS = CSH_ROWS * CSH_TEXELS_PER_ROW; + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + // Guards Csh_out against the previous iteration's readers. Inert on i == 0 + // but must stay unconditional to remain workgroup-uniform. + // coopmat-lds-fence: barrier() alone does NOT order shared stores against a + // subsequent coopMatLoad on the M51 Xclipse/AMD-PAL driver -- symptom is one + // stale MMA_M-row band of A, all columns, ~2.5% of runs, no crash (observed + // 2026-09-02 in sdpa_compute_out_coopmat.glsl). Measured cost of the fence: + // none (see this change's results). See memory + // `coopmat-lds-needs-explicit-memorybarriershared`. + memoryBarrierShared(); + barrier(); + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { +#ifdef HAS_BIAS + if (apply_bias > 0) { + const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopmat bias_tile; + coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); + result[i][j] += bias_tile; + } +#endif + coopmat out_tile = + coopmat(result[i][j]); + coopMatStore( + out_tile, Csh_out, + warpInTile.y * MMA_M * WG_TILE_N + + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j), + WG_TILE_N, + gl_CooperativeMatrixLayoutRowMajor); + } + memoryBarrierShared(); + barrier(); + + for (uint t = gl_LocalInvocationID.x; t < CSH_TEXELS; t += WG_SIZE) { + const uint lr = t / CSH_TEXELS_PER_ROW; + const uint lc4 = t % CSH_TEXELS_PER_ROW; + const uint m = + tile_m_start + (lr / MMA_M) * SG_TILE_M + i * MMA_M + (lr % MMA_M); + const uint base = lr * WG_TILE_N + lc4 * 4u; + imageStore( + t_output, + ivec3(tile_n_start / 4u + lc4, m, 0), + vec4( + float(Csh_out[base]), + float(Csh_out[base + 1u]), + float(Csh_out[base + 2u]), + float(Csh_out[base + 3u]))); + } + } +#else + const uint N_out = uint(out_N_arg); + [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { + [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { + const uint gi = tile_m_start + MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); + const uint gj = tile_n_start + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + +#ifdef HAS_BIAS + if (apply_bias > 0) { + const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); + coopmat bias_tile; + coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); + result[i][j] += bias_tile; + } +#endif + + coopmat out_tile = + coopmat(result[i][j]); + coopMatStore( + out_tile, t_output, + gi * N_out + gj, N_out, + gl_CooperativeMatrixLayoutRowMajor); + } + } +#endif // IO_TEXTURE +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr3.yaml b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr3.yaml new file mode 100644 index 00000000000..c4665a771d4 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr3.yaml @@ -0,0 +1,92 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# "zpgtr3": dbuf4zpgtr with tr3's SHARED-A layout. Three changes vs zpgtr, all +# on the A path only: +# 1. shared storage element type uint (4x int8 packed) -> scalar int8_t +# 2. slab-major (one MMA_K-wide slab per K step) -> row-major over the FULL +# chunk (WG_TILE_K wide) +# 3. A_ROW_PAD_I8 bytes of padding on each FULL-CHUNK row (tr3 uses 16) +# B staging, zp-hoist, byte-parallel nibble widening, the group/chunk nesting, +# the memoryBarrierShared fences and the fp16 output path are dbuf4zpgtr's, +# byte-for-byte unchanged. +# +# OPT-IN ONLY -- not a default. Select with +# ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4zpgtr3_txkgsp +# +# LDS cost of the A array is 2*WG_TILE_M*(WG_TILE_K+A_ROW_PAD_I8) bytes, i.e. +# padding is paid once per row, NOT once per MMA_K slab row. At +# t128x64k32g42s32 that is p0 8192 B / p8 10240 B / p16 12288 B against +# zpgtr's 8192 B, so p16 adds 4 KiB and is the only one of the three that can +# cost a workgroup of LDS-limited occupancy. + +linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr3: + parameter_names_with_default_values: + PRECISION: highp + HAS_BIAS: false + IO_STORAGE: buffer + WEIGHT_NBITS: 4 + WEIGHT_STORAGE: texture2d + MMA_M: 16 + MMA_N: 16 + MMA_K: 16 + WG_TILE_M: 64 + WG_TILE_N: 32 + WG_TILE_K: 32 + SG_GRID_X: 1 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 64 + A_ROW_PAD_I8: 16 + shader_variants: + # --- shipped geometry, pad swept: isolates padding from layout+element type + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr3_t128x64k32g42s32p16_buffer_texture2d_half + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + A_ROW_PAD_I8: 16 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr3_t128x64k32g42s32p8_buffer_texture2d_half + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + A_ROW_PAD_I8: 8 + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr3_t128x64k32g42s32p0_buffer_texture2d_half + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + A_ROW_PAD_I8: 0 + # --- iso geometry (the config TR3 was measured at) with tr3's own pad + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr3_t128x128k32g42s32p16_buffer_texture2d_half + WG_TILE_M: 128 + WG_TILE_N: 128 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + A_ROW_PAD_I8: 16 + # --- texture-IO coverage at the shipped geometry (storage-variant parity) + - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr3_t128x64k32g42s32p16_texture3d_texture2d_half + IO_STORAGE: texture3d + WG_TILE_M: 128 + WG_TILE_N: 64 + WG_TILE_K: 32 + SG_GRID_X: 4 + SG_GRID_Y: 2 + SUBGROUP_SIZE: 32 + A_ROW_PAD_I8: 16 + # NOTE on the INT8 weight path: the baseline dbuf4zpgtr yaml ships no + # WEIGHT_NBITS=8 variant either, and the host builds kernel names as + # ___half with no nbits field, so a w8 entry would be + # unreachable. The `#else` (INT8) branch of the shader is untouched by this + # change; it is simply not compiled, exactly as in the baseline. diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp index 3255622ba9e..4f6e79a2db3 100644 --- a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp @@ -128,6 +128,21 @@ static const char* const kDq8caTsweepPrefixes[] = { // staging/zp-hoist/nibble-widening unchanged). PROMOTED 2026-09-01 as // the shipped default -- see dq8ca_coopmat_variant() below. "tsweep_dbuf4zpgtr_t", + // dbuf4zpgtr with tr3's shared-A layout (scalar int8_t, row-major over + // the full chunk, A_ROW_PAD_I8 per-row pad). Correctness-clean 12/12 but + // MEASURED SLOWER (+4.71% best case) -- kept opt-in as a recorded + // negative result. Never a default. + "tsweep_dbuf4zpgtr3_t", + // B-staging ownership moved to one thread per 2-uint pair, all 256 threads + // active. bw2 keeps the uint array, bw3 retypes it to uvec2 for a real wide + // store. Both CORRECT (14/14 x3); neither faster (bw2 ~0%, bw3 +0.34%) -- + // the ISA shows scalar ds_store_b32 11->3 for nothing, because our LDS gap + // vs gemm-ubm TR3 is LOADS (80 vs 24), not stores (24 vs 8). OPT-IN only. + // (dbuf4zpgbv4/bw/bwr are KNOWN-INCORRECT -- deliberately NOT listed here + // so they cannot be selected; see their yaml headers for the isolated + // cause.) + "tsweep_dbuf4zpgbw2_t", + "tsweep_dbuf4zpgbw3_t", // (dq8ca-dequant-unpack-ablation Addendum 11 -- abl_aconst/abl_areadc/ // abl_abconst -- were measurement-only variants deleted once each // attribution was recorded; see openspec/changes/dq8ca-dequant-unpack- @@ -612,7 +627,12 @@ static bool dq8ca_variant_wants_rowmajor_a() { return v.rfind("tsweep_dbuf4tr_t", 0) == 0 || v.rfind("tsweep_dbuf4trm_t", 0) == 0 || v.rfind("tsweep_dbuf4trd_t", 0) == 0 || - v.rfind("tsweep_dbuf4zpgtr_t", 0) == 0; + v.rfind("tsweep_dbuf4zpgtr_t", 0) == 0 || + // zpgtr3 keeps zpgtr's GLOBAL A access verbatim; only the SHARED + // layout differs, so it has the same kPackedInt8_4W requirement. + v.rfind("tsweep_dbuf4zpgtr3_t", 0) == 0 || + v.rfind("tsweep_dbuf4zpgbw2_t", 0) == 0 || + v.rfind("tsweep_dbuf4zpgbw3_t", 0) == 0; } // Mirrors the coopmat branch of pick_linear_dqa_qw_shader() so graph-build time From 8ff908a365ae53bf8320752c1dceed2c282a950c Mon Sep 17 00:00:00 2001 From: Yanwen Xu Date: Tue, 15 Sep 2026 11:25:36 -0700 Subject: [PATCH 26/28] [ET-VK] Remove dead/superseded coopmat tsweep and SDPA-coop shader variants Housekeeping: prunes the tsweep GLSL/YAML files that no longer serve any purpose now that this branch ships a single validated default per family (q4gsw: tsweep_dbuf4_t128x128k16g22s32; dq8ca: tsweep_dbuf4zpgtr_t128x64k32g42s32) -- - dbuf1/dbuf2/dbuf3 and the bare "tsweep_t" prefixes: dead in both the q4gsw and dq8ca prefix lists (no GLSL ever existed for them in this branch), so selecting one could only ever crash confusingly at shader lookup instead of failing validation cleanly. - dq8ca dbuf4 / dbuf4zpg: the 2026-08-26 and 2026-08-28 defaults, each twice superseded. - dq8ca dbuf4zpgtr3: correctness-clean but measured +4.71% slower. - dq8ca dbuf4zpgbw2/bw3: correctness-clean, ~0% delta, no reason to keep as a build target. - sdpa_compute_attn_weights_coop / sdpa_compute_out_coop: unreferenced by any remaining code path. QuantizedLinear.cpp's prefix lists and dq8ca_coopmat_variant()'s invalid-token fallback are updated in lockstep so an env var can never name a shader that no longer exists. Full history (promotion dates, measured deltas, why each variant was rejected) is preserved in the comments below and in the deleted files themselves, on branch yanwen/release14-quant-shaders-archived-2026-09-15. Authored with Claude Code (Sonnet 5). Co-Authored-By: Claude Sonnet 5 --- ...near_dq8ca_q4gsw_coopmat_tsweep_dbuf4.glsl | 662 ----- ...near_dq8ca_q4gsw_coopmat_tsweep_dbuf4.yaml | 2188 ----------------- ...r_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpg.glsl | 760 ------ ...r_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpg.yaml | 80 - ...q8ca_q4gsw_coopmat_tsweep_dbuf4zpgbv4.glsl | 792 ------ ...q8ca_q4gsw_coopmat_tsweep_dbuf4zpgbv4.yaml | 68 - ...dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw.glsl | 781 ------ ...dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw.yaml | 68 - ...q8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw2.glsl | 762 ------ ...q8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw2.yaml | 60 - ...q8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw3.glsl | 798 ------ ...q8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw3.yaml | 60 - ...q8ca_q4gsw_coopmat_tsweep_dbuf4zpgbwr.glsl | 783 ------ ...q8ca_q4gsw_coopmat_tsweep_dbuf4zpgbwr.yaml | 68 - ...q8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr3.glsl | 781 ------ ...q8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr3.yaml | 92 - .../glsl/sdpa_compute_attn_weights_coop.glsl | 192 -- .../glsl/sdpa_compute_attn_weights_coop.yaml | 25 - .../graph/ops/glsl/sdpa_compute_out_coop.glsl | 199 -- .../graph/ops/glsl/sdpa_compute_out_coop.yaml | 25 - .../graph/ops/impl/QuantizedLinear.cpp | 89 +- 21 files changed, 24 insertions(+), 9309 deletions(-) delete mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4.glsl delete mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4.yaml delete mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpg.glsl delete mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpg.yaml delete mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbv4.glsl delete mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbv4.yaml delete mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw.glsl delete mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw.yaml delete mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw2.glsl delete mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw2.yaml delete mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw3.glsl delete mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw3.yaml delete mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbwr.glsl delete mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbwr.yaml delete mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr3.glsl delete mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr3.yaml delete mode 100644 backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_coop.glsl delete mode 100644 backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_coop.yaml delete mode 100644 backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coop.glsl delete mode 100644 backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coop.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4.glsl b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4.glsl deleted file mode 100644 index 80408194843..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4.glsl +++ /dev/null @@ -1,662 +0,0 @@ -/* - * Copyright (c) Meta Platforms, Inc. and affiliates. - * All rights reserved. - * - * This source code is licensed under the BSD-style license found in the - * LICENSE file in the root directory of this source tree. - */ - -/* - * TILE/SUBGROUP-SWEEP variant of the int8 dq8ca_q4gsw coopmat shader's dbuf4 - * ("store-first-for-next", the ORIGINAL loop structure before specs/025 User - * Story 1 picked dbuf2) loop structure (specs/041-dbuf4-tile-sweep). Forked - * from linear_dq8ca_q4gsw_coopmat_tsweep.glsl (which carries dbuf2's loop, - * the production winner) -- everything except the PROLOGUE/MAIN LOOP block - * is identical: bindings, spec-constants, tile-geometry templating, LDS - * layout (ColumnMajor B + skew), int8 WMMA thread maps, group epilog, - * bias/store epilogue. Only the loop structure is swapped to dbuf4, - * recovered from git commit 8d0f23ee78's - * linear_dq8ca_q4gsw_coopmat_dbuf4.glsl (see specs/041/reference/) -- the - * byte-identical pre-swap copy of what is now linear_dq8ca_qw_coopmat.glsl. - * - * The nested `groups x chunks` loop and unconditional group epilog are kept - * exactly as in dbuf2 -- flattening them crashes the Xclipse PAL compiler at - * large spec-resolved trip counts (see dbuf2's own header). Only the - * store/barrier/prefetch ORDER within each chunk iteration is inverted: - * - * dbuf2 (this file's base): store(temp, already prefetched -> cur slice) - * -> barrier -> MMA(cur) -> prefetch(next -> temp) [store owns the - * CURRENT chunk, at the iteration's start] - * dbuf4 (this file): barrier -> prefetch(next -> temp) -> MMA(cur) -> - * store(temp -> next slice) [store owns the NEXT chunk, at the - * iteration's end -- the mirror image] - * - * The group wsum/wsc ping-pong is inverted the same way: dbuf4 stores the - * next group's values (prefetched during the crossing chunk) at the TAIL of - * that chunk, instead of dbuf2's HEAD-of-new-group placement. - * - * Selected at dispatch via - * ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4_txkgs<32|64> - * (QuantizedLinear.cpp), additive to the existing tsweep_t... (dbuf2) - * namespace. - * - * KHR Cooperative Matrix variant of the dynamically-quantized-activation - * linear tiled shader (WEIGHT_NBITS=4): - * 4 -> linear_dq8ca_q4gsw_coopmat INT4 group-symmetric weight - * - * Performs: out[M,N] = dequant(int8_act) * dequant(int_w) (+ bias) - * via coopmat x coopmat -> coopmat on the matrix unit. - * - * Hard preconditions: - * M % WG_TILE_M == 0, N % WG_TILE_N == 0, K % WG_TILE_K == 0, - * INT4: group_size % WG_TILE_K == 0, - * device exposes coopmatx-> at 16x16x16. - */ - -#version 450 core - -#extension GL_KHR_cooperative_matrix : require -#extension GL_KHR_memory_scope_semantics : require -#extension GL_KHR_shader_subgroup_basic : enable -#extension GL_EXT_shader_explicit_arithmetic_types : require -#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require -#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require -#extension GL_EXT_control_flow_attributes : enable - -#define PRECISION ${PRECISION} - -$if WEIGHT_NBITS == 4: - #define WEIGHT_INT4 - -$if HAS_BIAS: - #define HAS_BIAS - -$if WEIGHT_STORAGE == "buffer": - #define WEIGHT_BUFFER - -$if IO_STORAGE == "texture3d": - #define IO_TEXTURE - -layout(std430) buffer; - -#include "common.glslh" - -// Bindings — match add_linear_dqa_qw_node arg order: -// output(0), fp_input(1), packed_int8_input(2), int_input_sums(3 - unused), -// input_scales(4), input_zps(5), packed_weight(6), weight_sums(7), -// weight_scales(8), bias(9). -${layout_declare_tensor(B, "w", "t_output", "half", IO_STORAGE, is_scalar_array=True)} -// t_input is unread here -- the activations arrive already quantized in -// t_packed_int8_input -- but stays declared so the binding layout matches the -// dispatch site. It tracks IO_STORAGE so the two IO tensors stay consistent. -${layout_declare_tensor(B, "r", "t_input", "half", IO_STORAGE, is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_packed_int8_input", "int", "buffer", is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_int8_input_sums", "int", "buffer", is_scalar_array=True)} -${layout_declare_tensor(B, "r", "t_int8_input_scales", "half", "texture3d")} -${layout_declare_tensor(B, "r", "t_int8_input_zps", "int8", "texture3d")} -${layout_declare_tensor(B, "r", "t_packed_weight", "int", WEIGHT_STORAGE, is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_weight_sums", "int", "buffer", is_scalar_array=True)} -${layout_declare_tensor(B, "r", "t_weight_scales", "half", "buffer", is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_bias", "half", "buffer", is_scalar_array=True)} - -${layout_declare_ubo(B, "ivec4", "output_sizes")} -${layout_declare_ubo(B, "ivec4", "input_sizes")} - -layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; - -${layout_declare_spec_const(C, "int", "apply_bias", "0")} -// INT4 only; inert (0) for INT8 so the dispatcher's spec list lines up. -${layout_declare_spec_const(C, "int", "K4_per_group", "0")} -${layout_declare_spec_const(C, "int", "num_groups_arg", "0")} -${layout_declare_spec_const(C, "int", "out_N_arg", "0")} - -// Tile geometry -const uint MMA_M = ${MMA_M}; -const uint MMA_N = ${MMA_N}; -const uint MMA_K = ${MMA_K}; - -const uint WG_TILE_M = ${WG_TILE_M}; -const uint WG_TILE_N = ${WG_TILE_N}; -const uint WG_TILE_K = ${WG_TILE_K}; - -const uint SG_GRID_X = ${SG_GRID_X}; -const uint SG_GRID_Y = ${SG_GRID_Y}; -const uint SUBGROUP_SIZE = ${SUBGROUP_SIZE}; -const uint NUM_SUBGROUPS = SG_GRID_X * SG_GRID_Y; -const uint WG_SIZE = NUM_SUBGROUPS * SUBGROUP_SIZE; - -const uint SG_TILE_M = WG_TILE_M / SG_GRID_Y; -const uint SG_TILE_N = WG_TILE_N / SG_GRID_X; -const uint MMAS_PER_SG_M = SG_TILE_M / MMA_M; -const uint MMAS_PER_SG_N = SG_TILE_N / MMA_N; - -const uint A_SLAB_INT8 = WG_TILE_M * MMA_K; -const uint B_USEFUL_U32 = MMA_K / 4u; -// No skew. The classic anti-bank-conflict "+1" padding was measured SLOWER on M51 at the shipped -// tile (t64x32k32g12s64): stride=4 (this) is +2.98pp efficiency over stride=5 (the old +1 skew), -// stride=6/+2 is a wash, stride=8/+4 is worse -- real, on-device, dq8ca-dequant-unpack-ablation -// (openspec/changes/dq8ca-dequant-unpack-ablation/results/README.md, Addendum 6). B_USEFUL_U32 is -// tile-invariant (MMA_K is fixed at 16 for every variant in this family), so this applies uniformly; -// only re-validated at the shipped tile specifically, not every swept variant. -const uint B_STRIDE_U32 = B_USEFUL_U32; -const uint B_SLAB_U32 = WG_TILE_N * B_STRIDE_U32; -const uint NUM_K_SLABS = WG_TILE_K / MMA_K; - -const uint A_SLAB_U32 = A_SLAB_INT8 / 4u; -const uint A_STRIDE_U32 = MMA_K / 4u; - -// One ping-pong slice covers all K-slabs of one chunk. -const uint ASH_SLICE_U32 = NUM_K_SLABS * A_SLAB_U32; -const uint BSH_SLICE_U32 = NUM_K_SLABS * B_SLAB_U32; - -// Double-buffered MMA operand staging. -shared uint Ash_int8[2u * ASH_SLICE_U32]; -shared uint Bsh_int8[2u * BSH_SLICE_U32]; - -// Per-WG-tile-row activation params (loaded ONCE at WG start; constant -// across groups). -shared int izp_sh[WG_TILE_M]; // int32 (cast from int8 source) for broadcast -shared float ifs_sh[WG_TILE_M]; // float32 (cast from fp16 source) for broadcast - -// Per-(group, output-channel) weight params, ping-ponged by group parity. -// (For per-channel INT8 only slice 0 is ever used.) -shared int wsum_sh[2u * WG_TILE_N]; -shared float wsc_sh[2u * WG_TILE_N]; - -#ifdef HAS_BIAS -shared float bias_sh[WG_TILE_N]; -#endif - -#ifdef IO_TEXTURE -// Result staging for the imageStore epilogue, mirroring the fp16 kernel: -// SG_GRID_Y bands of MMA_M rows, each WG_TILE_N wide, row-major. A full -// WG_TILE_M x WG_TILE_N buffer would cost SG_GRID_Y/MMAS_PER_SG_M x more LDS -// and wreck occupancy. float16_t-typed because coopMatStore needs it. -const uint CSH_ROWS = SG_GRID_Y * MMA_M; -shared float16_t Csh_out[CSH_ROWS * WG_TILE_N]; -#endif - -#ifdef WEIGHT_INT4 -// Coalesced-B-write address inversion: given the contiguous per-thread LDS -// index `a` (in [0, BSH_SLICE_U32)) this thread will write to, recover which -// global-fetch element it needs. Consecutive threads (consecutive `a`) now -// write consecutive LDS words -- unlike the pre-2026-08-26 mapping, where 8 -// consecutive threads wrote addresses B_STRIDE_U32 words apart. Real, -// on-device measurement (dq8ca-dequant-unpack-ablation follow-up, -// openspec/changes/dq8ca-dequant-unpack-ablation/results/README.md, "the -// coalesced B-store rewrite"): a consistent +0.7-0.8% real speedup across -// 1B/3B/8B at the shipped tile, correctness-clean (10/10 buffer + 10/10 -// texture3d, all three models). `chunkK_base` is the K-offset of the chunk -// being staged (0 for the prologue, chunkK_nxt for the main loop's chunk+1 -// prefetch; the store sites don't need it -- r/parity depend only on n_col, -// not chunkK_base -- so they pass 0u). -struct BCoalIndex { - uint n8_blk; - uint k4_blk; - uint r; - uint parity; -}; - -BCoalIndex bcoal_index(const uint a, const uint chunkK_base, const uint tile_n_start) { - const uint slab_idx = a / B_SLAB_U32; - const uint local_a = a % B_SLAB_U32; - const uint n_col = local_a / B_STRIDE_U32; - const uint k4_in_slab = local_a % B_STRIDE_U32; - const uint k4_in_chunk = slab_idx * (MMA_K >> 2u) + k4_in_slab; - const uint n8_in_tile = n_col >> 3u; - const uint rem = n_col & 7u; - BCoalIndex idx; - idx.n8_blk = (tile_n_start >> 3u) + n8_in_tile; - idx.k4_blk = (chunkK_base >> 2u) + k4_in_chunk; - idx.r = rem & 3u; - idx.parity = rem >> 2u; - return idx; -} -#endif - -// Running fp32 accumulator (across all groups). -coopmat - result[MMAS_PER_SG_M][MMAS_PER_SG_N]; - -// Per-group int32 MMA accumulator. -coopmat - accum_int32[MMAS_PER_SG_M][MMAS_PER_SG_N]; - -void main() { - const uvec2 tileID = uvec2(gl_WorkGroupID.xy); - const uvec2 warpInTile = uvec2( - gl_SubgroupID % SG_GRID_X, - gl_SubgroupID / SG_GRID_X); - - const uint K = uint(input_sizes.x); - const uint N = uint(output_sizes.x); - const uint N4 = (N + 3u) / 4u; - const uint nblocks_x_A = (K + 3u) >> 2u; - -#ifdef WEIGHT_INT4 - const uint num_groups = uint(num_groups_arg); - const uint CHUNKS_PER_GROUP = uint(K4_per_group) * 4u / WG_TILE_K; -#else - // Per-channel: a single quant "group" spanning all of K. The nested - // groups x chunks loop below collapses to a flat chunk loop, the wsum/wsc - // ping-pong never crosses a boundary, and the epilog runs exactly once. - const uint num_groups = 1u; - const uint CHUNKS_PER_GROUP = uint(num_groups_arg); -#endif - const uint num_chunks = num_groups * CHUNKS_PER_GROUP; - - const uint tile_m_start = WG_TILE_M * tileID.y; - const uint tile_n_start = WG_TILE_N * tileID.x; - - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - result[i][j] = coopmat(0.0); - accum_int32[i][j] = coopmat(0); - } - } - - // --- A staging thread map: one (m4, k4) ivec4 block per active thread --- - const uint K_BLOCKS_PER_CHUNK = WG_TILE_K >> 2u; - const uint A_ACTIVE_THREADS = (WG_TILE_M >> 2u) * K_BLOCKS_PER_CHUNK; - const uint a_m_block = gl_LocalInvocationID.x / K_BLOCKS_PER_CHUNK; - const uint a_k_block = gl_LocalInvocationID.x % K_BLOCKS_PER_CHUNK; - const bool a_active = gl_LocalInvocationID.x < A_ACTIVE_THREADS; - -#ifdef WEIGHT_INT4 - // --- B staging thread map: (block, col) slots; each slot extracts one - // ColumnMajor LDS uint (4 K-contiguous sign-extended int8) --- - const uint B_TOTAL_SLOTS = K_BLOCKS_PER_CHUNK * WG_TILE_N; - const uint B_SLOTS_PER_THREAD = B_TOTAL_SLOTS / WG_SIZE; - const uint N8_PER_TILE = WG_TILE_N >> 3u; -#else - // --- B staging thread map: one (k4, n4) ivec4 block per active thread --- - const uint B_FETCH_SLOTS = K_BLOCKS_PER_CHUNK * (WG_TILE_N >> 2u); - const uint N4_PER_TILE = WG_TILE_N >> 2u; - const uint b_k4_in_chunk = gl_LocalInvocationID.x / N4_PER_TILE; - const uint b_n_uint_col = gl_LocalInvocationID.x % N4_PER_TILE; - const bool b_active = gl_LocalInvocationID.x < B_FETCH_SLOTS; -#endif - - // Prefetch temp registers. - ivec4 temp_A; -#ifdef WEIGHT_INT4 - ivec4 temp_B[B_SLOTS_PER_THREAD]; - int temp_wsum; - float temp_wsc; -#else - ivec4 temp_B; -#endif - - // ========================================================= - // PROLOGUE - // ========================================================= - if (gl_LocalInvocationID.x < (WG_TILE_M >> 2u)) { - const uint m4 = (tile_m_start >> 2u) + gl_LocalInvocationID.x; - const vec4 sc = vec4(texelFetch(t_int8_input_scales, ivec3(m4, 0, 0), 0)); - const ivec4 zp = texelFetch(t_int8_input_zps, ivec3(m4, 0, 0), 0); - const uint base = gl_LocalInvocationID.x * 4u; - ifs_sh[base + 0u] = sc.x; ifs_sh[base + 1u] = sc.y; - ifs_sh[base + 2u] = sc.z; ifs_sh[base + 3u] = sc.w; - izp_sh[base + 0u] = zp.x; izp_sh[base + 1u] = zp.y; - izp_sh[base + 2u] = zp.z; izp_sh[base + 3u] = zp.w; - } - // Group 0 weight sums/scales -> slice 0. - if (gl_LocalInvocationID.x < WG_TILE_N) { - const uint n_idx = tile_n_start + gl_LocalInvocationID.x; - f16vec4 sv = t_weight_scales[n_idx >> 2u]; - wsc_sh[gl_LocalInvocationID.x] = float(sv[n_idx & 3u]); - wsum_sh[gl_LocalInvocationID.x] = t_weight_sums[n_idx]; - } - memoryBarrierShared(); - barrier(); - - // izp/ifs are per-row activation params, constant across K groups — - // broadcast them into coopmats ONCE; the group epilog reuses them every - // group (they depend only on the row block i, not on the group or j). - coopmat - izp_bcast[MMAS_PER_SG_M]; - coopmat - ifs_bcast[MMAS_PER_SG_M]; - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - const uint local_m_base = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); - coopMatLoad( - izp_bcast[i], izp_sh, - local_m_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutColumnMajor); - coopMatLoad( - ifs_bcast[i], ifs_sh, - local_m_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutColumnMajor); - } - - // dbuf4: prefetch chunk 0 into temp registers, THEN store to slice 0 (no - // barrier here -- the main loop's first iteration barriers before - // reading slice 0). - if (a_active) { - const uint m4_global = (tile_m_start >> 2u) + a_m_block; - temp_A = t_packed_int8_input[m4_global * nblocks_x_A + a_k_block]; - } -#ifdef WEIGHT_INT4 - [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { - const uint a = gl_LocalInvocationID.x + si * WG_SIZE; - const BCoalIndex bidx0 = bcoal_index(a, 0u, tile_n_start); -#ifdef WEIGHT_BUFFER - temp_B[si] = t_packed_weight[(bidx0.n8_blk * nblocks_x_A) + bidx0.k4_blk]; -#else - temp_B[si] = texelFetch(t_packed_weight, ivec2(bidx0.k4_blk, bidx0.n8_blk), 0); -#endif - } -#else - if (b_active) { - const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; -#ifdef WEIGHT_BUFFER - temp_B = t_packed_weight[(b_k4_in_chunk * N4) + block_x_w]; -#else - temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, b_k4_in_chunk), 0); -#endif - } -#endif - { - // store chunk 0 -> slice 0 - if (a_active) { - const uint slab_idx = a_k_block / (MMA_K >> 2u); - const uint k_uint_in_slab = a_k_block % (MMA_K >> 2u); - const uint base_row = a_m_block * 4u; - [[unroll]] for (uint m4i = 0; m4i < 4u; ++m4i) { - Ash_int8[slab_idx * A_SLAB_U32 + (base_row + m4i) * A_STRIDE_U32 + k_uint_in_slab] = - uint(temp_A[m4i]); - } - } -#ifdef WEIGHT_INT4 - [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { - const uint a = gl_LocalInvocationID.x + si * WG_SIZE; - const BCoalIndex bidx0 = bcoal_index(a, 0u, tile_n_start); - const int w = temp_B[si][bidx0.r]; - const int base = int(4u * bidx0.parity); - const int v0 = (((w >> (base + 0)) & 0xF) - 8) & 0xFF; - const int v1 = (((w >> (base + 8)) & 0xF) - 8) & 0xFF; - const int v2 = (((w >> (base + 16)) & 0xF) - 8) & 0xFF; - const int v3 = (((w >> (base + 24)) & 0xF) - 8) & 0xFF; - Bsh_int8[a] = - uint(v0 | (v1 << 8) | (v2 << 16) | (v3 << 24)); - } -#else - if (b_active) { - const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); - const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); - const uint n_col_base = b_n_uint_col * 4u; - [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { - Bsh_int8[slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab] = - uint(temp_B[n_in_blk]); - } - } -#endif - } - - // ========================================================= - // MAIN LOOP (dbuf4) — nested groups x chunks (kept nested; flattening it - // with a conditional coopmat epilog crashes the Xclipse PAL compiler at - // large spec-resolved trip counts). One barrier per chunk. Chunk - // iteration (global index `chunk`): - // 1. barrier — A/B slice (chunk%2) fully written; on the first chunk - // of group g, wsum/wsc slice (g%2) is too. - // 2. prefetch — chunk+1 (A blocks, B blocks) into temp; when chunk+1 - // starts a new group, also its wsum/wsc element. Skipped - // entirely on the final chunk. - // 3. int8 MMA — on slice (chunk%2) into accum_int32. - // 4. store — temp -> A/B slice ((chunk+1)%2), unpacking the weight; - // on a group boundary, wsum/wsc -> slice ((g+1)%2). - // The group epilog runs unconditionally at the tail of each group. - // ========================================================= - uint chunk = 0; - for (uint group_i = 0; group_i < num_groups; ++group_i) { - for (uint inner = 0; inner < CHUNKS_PER_GROUP; ++inner, ++chunk) { - const bool has_next = chunk + 1u < num_chunks; - const bool group_crossing = has_next && (inner + 1u == CHUNKS_PER_GROUP); - const uint cur_a = (chunk % 2u) * ASH_SLICE_U32; - const uint cur_b = (chunk % 2u) * BSH_SLICE_U32; - const uint nxt_a = ((chunk + 1u) % 2u) * ASH_SLICE_U32; - const uint nxt_b = ((chunk + 1u) % 2u) * BSH_SLICE_U32; - - barrier(); - - // --- 2. prefetch chunk+1 -> temp --- - if (has_next) { - const uint chunkK_nxt = (chunk + 1u) * WG_TILE_K; - if (a_active) { - const uint m4_global = (tile_m_start >> 2u) + a_m_block; - const uint k4_global = (chunkK_nxt >> 2u) + a_k_block; - temp_A = t_packed_int8_input[m4_global * nblocks_x_A + k4_global]; - } -#ifdef WEIGHT_INT4 - [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { - const uint a = gl_LocalInvocationID.x + si * WG_SIZE; - const BCoalIndex bidx1 = bcoal_index(a, chunkK_nxt, tile_n_start); -#ifdef WEIGHT_BUFFER - temp_B[si] = t_packed_weight[(bidx1.n8_blk * nblocks_x_A) + bidx1.k4_blk]; -#else - temp_B[si] = texelFetch(t_packed_weight, ivec2(bidx1.k4_blk, bidx1.n8_blk), 0); -#endif - } - if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { - const uint n_idx = tile_n_start + gl_LocalInvocationID.x; - f16vec4 sv = t_weight_scales[(group_i + 1u) * N4 + (n_idx >> 2u)]; - temp_wsc = float(sv[n_idx & 3u]); - temp_wsum = t_weight_sums[(group_i + 1u) * N + n_idx]; - } -#else - if (b_active) { - const uint block_y_w = (chunkK_nxt >> 2u) + b_k4_in_chunk; - const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; -#ifdef WEIGHT_BUFFER - temp_B = t_packed_weight[(block_y_w * N4) + block_x_w]; -#else - temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, block_y_w), 0); -#endif - } -#endif - } - - // --- 3. int8 MMA on the cur slice --- - [[unroll]] for (uint k = 0; k < NUM_K_SLABS; ++k) { - const uint slab_a_base_u32 = cur_a + k * A_SLAB_U32; - const uint slab_b_base_u32 = cur_b + k * B_SLAB_U32; - - coopmat matA[MMAS_PER_SG_M]; - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - const uint row_a = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); - coopMatLoad( - matA[i], Ash_int8, - slab_a_base_u32 + row_a * A_STRIDE_U32, - A_STRIDE_U32, - gl_CooperativeMatrixLayoutRowMajor); - } - - coopmat matB; - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint col_b = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopMatLoad( - matB, Bsh_int8, - slab_b_base_u32 + col_b * B_STRIDE_U32, - B_STRIDE_U32, - gl_CooperativeMatrixLayoutColumnMajor); - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - accum_int32[i][j] = coopMatMulAdd(matA[i], matB, accum_int32[i][j]); - } - } - } - - // --- 4. store temp (chunk+1) -> nxt slice --- - if (has_next) { - if (a_active) { - const uint slab_idx = a_k_block / (MMA_K >> 2u); - const uint k_uint_in_slab = a_k_block % (MMA_K >> 2u); - const uint base_row = a_m_block * 4u; - [[unroll]] for (uint m4i = 0; m4i < 4u; ++m4i) { - Ash_int8[nxt_a + slab_idx * A_SLAB_U32 + (base_row + m4i) * A_STRIDE_U32 + k_uint_in_slab] = - uint(temp_A[m4i]); - } - } -#ifdef WEIGHT_INT4 - [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { - const uint a = gl_LocalInvocationID.x + si * WG_SIZE; - // r/parity depend only on n_col (a % B_SLAB_U32, mod 8) -- not on - // chunkK_base -- so this store site can pass a dummy 0u (unlike - // the fetch site above, a separate `if (has_next)` scope, which - // needs the real chunkK_nxt to compute n8_blk/k4_blk). - const BCoalIndex bidx1 = bcoal_index(a, 0u, tile_n_start); - const int w = temp_B[si][bidx1.r]; - const int base = int(4u * bidx1.parity); - const int v0 = (((w >> (base + 0)) & 0xF) - 8) & 0xFF; - const int v1 = (((w >> (base + 8)) & 0xF) - 8) & 0xFF; - const int v2 = (((w >> (base + 16)) & 0xF) - 8) & 0xFF; - const int v3 = (((w >> (base + 24)) & 0xF) - 8) & 0xFF; - Bsh_int8[nxt_b + a] = - uint(v0 | (v1 << 8) | (v2 << 16) | (v3 << 24)); - } - if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { - const uint wbase_nxt = ((group_i + 1u) % 2u) * WG_TILE_N; - wsum_sh[wbase_nxt + gl_LocalInvocationID.x] = temp_wsum; - wsc_sh[wbase_nxt + gl_LocalInvocationID.x] = temp_wsc; - } -#else - if (b_active) { - const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); - const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); - const uint n_col_base = b_n_uint_col * 4u; - [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { - Bsh_int8[nxt_b + slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab] = - uint(temp_B[n_in_blk]); - } - } -#endif - } - } // chunks - - // --- Group epilog: dequant accum_int32 -> result, reset accum --- - { - const uint wbase = (group_i % 2u) * WG_TILE_N; - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - - coopmat wsum_bcast; - coopMatLoad( - wsum_bcast, wsum_sh, - wbase + local_n_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutRowMajor); - - coopmat wsc_bcast; - coopMatLoad( - wsc_bcast, wsc_sh, - wbase + local_n_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutRowMajor); - - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - coopmat adjusted = - accum_int32[i][j] - izp_bcast[i] * wsum_bcast; - coopmat adjusted_fp = - coopmat(adjusted); - coopmat scales_outer = - ifs_bcast[i] * wsc_bcast; - result[i][j] += adjusted_fp * scales_outer; - accum_int32[i][j] = coopmat(0); - } - } - } - } // groups - - // --- Bias (optional) --- -#ifdef HAS_BIAS - if (apply_bias > 0) { - for (uint t = gl_LocalInvocationID.x; t < WG_TILE_N; t += WG_SIZE) { - bias_sh[t] = float(t_bias[tile_n_start + t]); - } - memoryBarrierShared(); - barrier(); - } -#endif - - // --- Store result tile --- - // N for the store address math MUST come from the spec constant, not the - // sizes UBO (see out_N_arg above). -#ifdef IO_TEXTURE - // Epilogue iteration i drains accumulator row-block i from EVERY subgroup - // into Csh_out at once, so the SG_GRID_Y bands it holds are disjoint global - // row ranges; the whole workgroup then imageStores them. lr / MMA_M is the - // writing subgroup's warpInTile.y, so the global row reproduces the buffer - // path's gi exactly. - // - // PORTABILITY NOTE: the barrier() in the loop body keeps this loop rolled - // despite [[unroll]], so result[i][j] IS dynamically indexed. Coopmat arrays - // are opaque per-lane storage and dynamic indexing is exactly the construct - // the Xclipse/AMD-PAL compiler has broken before -- check this first if the - // texture variants miscompile on M51. - const uint CSH_TEXELS_PER_ROW = WG_TILE_N / 4u; - const uint CSH_TEXELS = CSH_ROWS * CSH_TEXELS_PER_ROW; - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - // Guards Csh_out against the previous iteration's readers. Inert on i == 0 - // but must stay unconditional to remain workgroup-uniform. - barrier(); - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { -#ifdef HAS_BIAS - if (apply_bias > 0) { - const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopmat bias_tile; - coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); - result[i][j] += bias_tile; - } -#endif - coopmat out_tile = - coopmat(result[i][j]); - coopMatStore( - out_tile, Csh_out, - warpInTile.y * MMA_M * WG_TILE_N + - MMA_N * (MMAS_PER_SG_N * warpInTile.x + j), - WG_TILE_N, - gl_CooperativeMatrixLayoutRowMajor); - } - memoryBarrierShared(); - barrier(); - - for (uint t = gl_LocalInvocationID.x; t < CSH_TEXELS; t += WG_SIZE) { - const uint lr = t / CSH_TEXELS_PER_ROW; - const uint lc4 = t % CSH_TEXELS_PER_ROW; - const uint m = - tile_m_start + (lr / MMA_M) * SG_TILE_M + i * MMA_M + (lr % MMA_M); - const uint base = lr * WG_TILE_N + lc4 * 4u; - imageStore( - t_output, - ivec3(tile_n_start / 4u + lc4, m, 0), - vec4( - float(Csh_out[base]), - float(Csh_out[base + 1u]), - float(Csh_out[base + 2u]), - float(Csh_out[base + 3u]))); - } - } -#else - const uint N_out = uint(out_N_arg); - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint gi = tile_m_start + MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); - const uint gj = tile_n_start + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - -#ifdef HAS_BIAS - if (apply_bias > 0) { - const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopmat bias_tile; - coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); - result[i][j] += bias_tile; - } -#endif - - coopmat out_tile = - coopmat(result[i][j]); - coopMatStore( - out_tile, t_output, - gi * N_out + gj, N_out, - gl_CooperativeMatrixLayoutRowMajor); - } - } -#endif // IO_TEXTURE -} diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4.yaml b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4.yaml deleted file mode 100644 index 299ba97977d..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4.yaml +++ /dev/null @@ -1,2188 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -# specs/041-dbuf4-tile-sweep: TILE/SUBGROUP SWEEP variants of the int8 -# dq8ca_q4gsw coopmat shader's dbuf4 loop structure (the ORIGINAL loop before -# specs/025 User Story 1 picked dbuf2) -- -# linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4.glsl is a fork of -# linear_dq8ca_q4gsw_coopmat_tsweep.glsl (dbuf2) with only the loop structure -# swapped. Only tile geometry (WG_TILE_*, SG_GRID_*, SUBGROUP_SIZE) varies -# per variant. Selected via -# ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4_txkgs. -# Seed variant below matches the current production 8da4w tile (dbuf2, -# specs/026/027) as a legal, known-fast starting point for sweep.py's -# Optuna search; specs/041's sweep appends further candidates here as it -# runs. - -linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4: - parameter_names_with_default_values: - PRECISION: highp - HAS_BIAS: false - IO_STORAGE: buffer - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - MMA_M: 16 - MMA_N: 16 - MMA_K: 16 - WG_TILE_M: 64 - WG_TILE_N: 32 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - shader_variants: - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x32k32g12s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 32 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x32k32g12s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 32 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x256k16g82s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 256 - WG_TILE_K: 16 - SG_GRID_X: 8 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x256k16g82s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 256 - WG_TILE_K: 16 - SG_GRID_X: 8 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t256x128k64g44s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 256 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 4 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t256x128k64g44s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 256 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 4 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x64k128g44s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 64 - WG_TILE_K: 128 - SG_GRID_X: 4 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x64k128g44s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 64 - WG_TILE_K: 128 - SG_GRID_X: 4 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t256x32k64g21s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 256 - WG_TILE_N: 32 - WG_TILE_K: 64 - SG_GRID_X: 2 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t256x32k64g21s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 256 - WG_TILE_N: 32 - WG_TILE_K: 64 - SG_GRID_X: 2 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k64g42s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 64 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k64g42s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 64 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k16g41s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 4 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k16g41s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 4 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k128g84s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 128 - SG_GRID_X: 8 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k128g84s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 128 - SG_GRID_X: 8 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k16g24s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k16g24s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x256k32g24s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 256 - WG_TILE_K: 32 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x256k32g24s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 256 - WG_TILE_K: 32 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x32k64g11s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 32 - WG_TILE_K: 64 - SG_GRID_X: 1 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x32k64g11s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 32 - WG_TILE_K: 64 - SG_GRID_X: 1 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t16x128k64g21s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 16 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 2 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t16x128k64g21s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 16 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 2 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k64g81s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 8 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k64g81s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 8 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t32x256k64g22s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 32 - WG_TILE_N: 256 - WG_TILE_K: 64 - SG_GRID_X: 2 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t32x256k64g22s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 32 - WG_TILE_N: 256 - WG_TILE_K: 64 - SG_GRID_X: 2 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t256x128k32g12s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 256 - WG_TILE_N: 128 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t256x128k32g12s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 256 - WG_TILE_N: 128 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t256x64k16g14s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 256 - WG_TILE_N: 64 - WG_TILE_K: 16 - SG_GRID_X: 1 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t256x64k16g14s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 256 - WG_TILE_N: 64 - WG_TILE_K: 16 - SG_GRID_X: 1 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x16k32g12s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 16 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x16k32g12s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 16 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k16g24s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k16g24s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t32x32k32g12s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 32 - WG_TILE_N: 32 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t32x32k32g12s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 32 - WG_TILE_N: 32 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k64g22s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 64 - SG_GRID_X: 2 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k64g22s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 64 - SG_GRID_X: 2 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k16g24s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k16g24s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x32k16g12s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 32 - WG_TILE_K: 16 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x32k16g12s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 32 - WG_TILE_K: 16 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t32x128k64g11s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 32 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 1 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t32x128k64g11s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 32 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 1 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k32g41s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 4 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k32g41s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 4 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t256x128k16g81s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 256 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 8 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t256x128k16g81s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 256 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 8 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k16g28s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 2 - SG_GRID_Y: 8 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k16g28s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 2 - SG_GRID_Y: 8 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k32g44s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 32 - SG_GRID_X: 4 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k32g44s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 32 - SG_GRID_X: 4 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x256k16g81s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 256 - WG_TILE_K: 16 - SG_GRID_X: 8 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x256k16g81s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 256 - WG_TILE_K: 16 - SG_GRID_X: 8 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k16g12s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k16g12s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k32g11s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k32g11s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t256x128k64g82s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 256 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 8 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t256x128k64g82s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 256 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 8 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x32k128g11s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 32 - WG_TILE_K: 128 - SG_GRID_X: 1 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x32k128g11s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 32 - WG_TILE_K: 128 - SG_GRID_X: 1 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x32k64g14s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 32 - WG_TILE_K: 64 - SG_GRID_X: 1 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x32k64g14s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 32 - WG_TILE_K: 64 - SG_GRID_X: 1 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k64g24s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k64g24s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x256k16g24s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 256 - WG_TILE_K: 16 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x256k16g24s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 256 - WG_TILE_K: 16 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k64g28s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 2 - SG_GRID_Y: 8 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k64g28s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 2 - SG_GRID_Y: 8 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k64g18s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 64 - SG_GRID_X: 1 - SG_GRID_Y: 8 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k64g18s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 64 - SG_GRID_X: 1 - SG_GRID_Y: 8 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k64g42s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k64g42s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x64k16g22s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 64 - WG_TILE_K: 16 - SG_GRID_X: 2 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x64k16g22s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 64 - WG_TILE_K: 16 - SG_GRID_X: 2 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t32x256k16g12s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 32 - WG_TILE_N: 256 - WG_TILE_K: 16 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t32x256k16g12s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 32 - WG_TILE_N: 256 - WG_TILE_K: 16 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x256k16g28s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 256 - WG_TILE_K: 16 - SG_GRID_X: 2 - SG_GRID_Y: 8 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x256k16g28s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 256 - WG_TILE_K: 16 - SG_GRID_X: 2 - SG_GRID_Y: 8 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t16x16k32g11s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 16 - WG_TILE_N: 16 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t16x16k32g11s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 16 - WG_TILE_N: 16 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k32g14s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k32g14s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k32g12s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k32g12s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x256k32g12s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 256 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x256k32g12s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 256 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t256x128k16g22s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 256 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 2 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t256x128k16g22s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 256 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 2 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t256x32k64g12s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 256 - WG_TILE_N: 32 - WG_TILE_K: 64 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t256x32k64g12s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 256 - WG_TILE_N: 32 - WG_TILE_K: 64 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k128g42s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 128 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k128g42s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 128 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x256k64g18s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 256 - WG_TILE_K: 64 - SG_GRID_X: 1 - SG_GRID_Y: 8 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x256k64g18s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 256 - WG_TILE_K: 64 - SG_GRID_X: 1 - SG_GRID_Y: 8 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k64g14s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 1 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k64g14s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 1 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x32k64g24s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 32 - WG_TILE_K: 64 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x32k64g24s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 32 - WG_TILE_K: 64 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k64g81s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 8 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k64g81s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 8 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k128g22s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 128 - SG_GRID_X: 2 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k128g22s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 128 - SG_GRID_X: 2 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x64k64g12s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 64 - WG_TILE_K: 64 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x64k64g12s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 64 - WG_TILE_K: 64 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x256k64g82s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 256 - WG_TILE_K: 64 - SG_GRID_X: 8 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x256k64g82s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 256 - WG_TILE_K: 64 - SG_GRID_X: 8 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t32x64k64g42s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 32 - WG_TILE_N: 64 - WG_TILE_K: 64 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t32x64k64g42s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 32 - WG_TILE_N: 64 - WG_TILE_K: 64 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x32k16g12s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 32 - WG_TILE_K: 16 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x32k16g12s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 32 - WG_TILE_K: 16 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k16g42s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k16g42s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t256x128k32g48s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 256 - WG_TILE_N: 128 - WG_TILE_K: 32 - SG_GRID_X: 4 - SG_GRID_Y: 8 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t256x128k32g48s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 256 - WG_TILE_N: 128 - WG_TILE_K: 32 - SG_GRID_X: 4 - SG_GRID_Y: 8 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k32g44s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 32 - SG_GRID_X: 4 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k32g44s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 32 - SG_GRID_X: 4 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k16g18s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 1 - SG_GRID_Y: 8 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k16g18s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 1 - SG_GRID_Y: 8 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t32x128k16g21s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 32 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 2 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t32x128k16g21s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 32 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 2 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t16x128k64g81s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 16 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 8 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t16x128k64g81s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 16 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 8 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k16g18s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 16 - SG_GRID_X: 1 - SG_GRID_Y: 8 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k16g18s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 16 - SG_GRID_X: 1 - SG_GRID_Y: 8 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k64g22s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 64 - SG_GRID_X: 2 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k64g22s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 64 - SG_GRID_X: 2 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k32g24s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 32 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k32g24s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 32 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x32k64g22s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 32 - WG_TILE_K: 64 - SG_GRID_X: 2 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x32k64g22s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 32 - WG_TILE_K: 64 - SG_GRID_X: 2 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x32k64g28s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 32 - WG_TILE_K: 64 - SG_GRID_X: 2 - SG_GRID_Y: 8 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x32k64g28s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 32 - WG_TILE_K: 64 - SG_GRID_X: 2 - SG_GRID_Y: 8 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k32g41s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 32 - SG_GRID_X: 4 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k32g41s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 32 - SG_GRID_X: 4 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k128g28s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 128 - SG_GRID_X: 2 - SG_GRID_Y: 8 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k128g28s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 128 - SG_GRID_X: 2 - SG_GRID_Y: 8 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t256x32k16g11s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 256 - WG_TILE_N: 32 - WG_TILE_K: 16 - SG_GRID_X: 1 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t256x32k16g11s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 256 - WG_TILE_N: 32 - WG_TILE_K: 16 - SG_GRID_X: 1 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k64g44s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 4 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k64g44s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 4 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k16g14s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 16 - SG_GRID_X: 1 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k16g14s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 16 - SG_GRID_X: 1 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x16k64g12s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 16 - WG_TILE_K: 64 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x16k64g12s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 16 - WG_TILE_K: 64 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k16g21s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 16 - SG_GRID_X: 2 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k16g21s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 16 - SG_GRID_X: 2 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k32g18s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 8 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k32g18s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 8 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k64g24s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k64g24s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k32g42s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 32 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k32g42s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 32 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k128g18s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 128 - SG_GRID_X: 1 - SG_GRID_Y: 8 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k128g18s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 128 - SG_GRID_X: 1 - SG_GRID_Y: 8 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x16k64g18s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 16 - WG_TILE_K: 64 - SG_GRID_X: 1 - SG_GRID_Y: 8 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x16k64g18s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 16 - WG_TILE_K: 64 - SG_GRID_X: 1 - SG_GRID_Y: 8 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k64g21s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 2 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k64g21s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 2 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k32g24s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 32 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k32g24s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 32 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x256k32g84s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 256 - WG_TILE_K: 32 - SG_GRID_X: 8 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x256k32g84s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 256 - WG_TILE_K: 32 - SG_GRID_X: 8 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k16g44s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 4 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k16g44s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 4 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k32g12s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k32g12s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k64g11s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 64 - SG_GRID_X: 1 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k64g11s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 64 - SG_GRID_X: 1 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x16k64g12s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 16 - WG_TILE_K: 64 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x16k64g12s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 16 - WG_TILE_K: 64 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t256x128k64g24s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 256 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t256x128k64g24s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 256 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k32g24s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k32g24s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x32k128g12s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 32 - WG_TILE_K: 128 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x32k128g12s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 32 - WG_TILE_K: 128 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k128g44s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 128 - SG_GRID_X: 4 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k128g44s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 128 - SG_GRID_X: 4 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k32g24s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 32 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k32g24s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 32 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k64g84s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 8 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k64g84s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 8 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t32x16k32g12s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 32 - WG_TILE_N: 16 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t32x16k32g12s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 32 - WG_TILE_N: 16 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x64k32g42s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x64k32g42s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t256x128k64g81s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 256 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 8 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t256x128k64g81s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 256 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 8 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k64g44s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 64 - SG_GRID_X: 4 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k64g44s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 64 - SG_GRID_X: 4 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x64k32g24s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x64k32g24s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x256k16g14s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 256 - WG_TILE_K: 16 - SG_GRID_X: 1 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x256k16g14s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 256 - WG_TILE_K: 16 - SG_GRID_X: 1 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x256k32g44s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 256 - WG_TILE_K: 32 - SG_GRID_X: 4 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x256k32g44s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 256 - WG_TILE_K: 32 - SG_GRID_X: 4 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x32k32g24s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 32 - WG_TILE_K: 32 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x32k32g24s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 32 - WG_TILE_K: 32 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k64g24s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 64 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k64g24s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 64 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k32g14s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k32g14s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x32k16g12s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 32 - WG_TILE_K: 16 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x32k16g12s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 32 - WG_TILE_K: 16 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k32g14s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k32g14s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k32g24s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 32 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x128k32g24s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 32 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k16g24s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 16 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x64k16g24s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 16 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x256k16g12s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 256 - WG_TILE_K: 16 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t64x256k16g12s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 256 - WG_TILE_K: 16 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x256k16g24s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 256 - WG_TILE_K: 16 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x256k16g24s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 256 - WG_TILE_K: 16 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k64g28s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 2 - SG_GRID_Y: 8 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k64g28s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 2 - SG_GRID_Y: 8 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k16g14s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 1 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k16g14s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 1 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 64 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x256k32g18s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 256 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 8 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x256k32g18s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 256 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 8 - SUBGROUP_SIZE: 32 - # Reference-matching tile: shmem_double_buf4-tr.comp's own geometry - # (TILE 128x128x16, WORKGROUP_WIDTH_IN_SUBGROUPS=4 x DBUF4_GRID_HEIGHT=2, - # wave32) => 8 subgroups / 256 threads, MMAS_PER_SG 4x2. Added so our - # kernel can be A/B'd against the teammate's PAL capture at the same - # tile and the same 2048x1024x4096 shape (8b wk_wv prefill). - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k16g42s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4_t128x128k16g42s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpg.glsl b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpg.glsl deleted file mode 100644 index d5af04c92c7..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpg.glsl +++ /dev/null @@ -1,760 +0,0 @@ -/* - * Copyright (c) Meta Platforms, Inc. and affiliates. - * All rights reserved. - * - * This source code is licensed under the BSD-style license found in the - * LICENSE file in the root directory of this source tree. - */ - -/* - * "zp-hoisted" variant: identical to dbuf4 except that the activation - * zero-point correction and the per-row activation scale are applied ONCE - * after the group loop instead of once per quantization group. - * - * The per-group epilogue term factors exactly: - * - * out[m][n] = ifs[m] * SUM_g wsc[g][n] * ( acc[g] - izp[m]*wsum[g][n] ) - * = ifs[m] * [ SUM_g wsc[g][n]*acc[g] - izp[m]*SUM_g wsc[g][n]*wsum[g][n] ] - * \__ weight-side only __/ - * - * `ifs` is per-row and group-independent so it factors out entirely, and the - * zero-point term separates into a per-row scalar times a per-output-channel - * weight-side sum. That sum depends on no activation data, so it is - * accumulated once into `wcorr_sh` in the prologue rather than being rebuilt - * per group. - * - * Consequences vs dbuf4, per accumulator tile per group: - * - gone: izp*wsum multiply and the subtract (48 v_sub* in the loop body) - * - gone: ifs*wsc multiply (part of 48 dequant-fp) - * - gone: the wsum_sh shared array and its ping-pong - * - gone: izp_bcast / ifs_bcast live across the loop (register pressure) - * - kept: result += float(acc) * wsc - * - * Exact in exact arithmetic, but NOT bit-exact in fp32 -- the summation order - * changes -- so it is gated on the correctness matrix like any other change. - * - * No new binding and no export-format change: the weight-side sum is derived - * in the prologue from t_weight_scales and t_weight_sums, both already bound. - * - * INTERVENTION F variant: all loop-invariant staging index arithmetic is - * hoisted out of the group loop (see "INTERVENTION F" below). - * Additionally widens int4 -> int8 byte-parallel (see widen_nibbles below), - * replacing the per-nibble shift/mask/bias-subtract chain. Bit-identical. - * - * Selected via - * ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4zpn_txkgs<32|64>. - * - * Original dbuf4 header follows. - * - * TILE/SUBGROUP-SWEEP variant of the int8 dq8ca_q4gsw coopmat shader's dbuf4 - * ("store-first-for-next", the ORIGINAL loop structure before specs/025 User - * Story 1 picked dbuf2) loop structure (specs/041-dbuf4-tile-sweep). Forked - * from linear_dq8ca_q4gsw_coopmat_tsweep.glsl (which carries dbuf2's loop, - * the production winner) -- everything except the PROLOGUE/MAIN LOOP block - * is identical: bindings, spec-constants, tile-geometry templating, LDS - * layout (ColumnMajor B + skew), int8 WMMA thread maps, group epilog, - * bias/store epilogue. Only the loop structure is swapped to dbuf4, - * recovered from git commit 8d0f23ee78's - * linear_dq8ca_q4gsw_coopmat_dbuf4.glsl (see specs/041/reference/) -- the - * byte-identical pre-swap copy of what is now linear_dq8ca_qw_coopmat.glsl. - * - * The nested `groups x chunks` loop and unconditional group epilog are kept - * exactly as in dbuf2 -- flattening them crashes the Xclipse PAL compiler at - * large spec-resolved trip counts (see dbuf2's own header). Only the - * store/barrier/prefetch ORDER within each chunk iteration is inverted: - * - * dbuf2 (this file's base): store(temp, already prefetched -> cur slice) - * -> barrier -> MMA(cur) -> prefetch(next -> temp) [store owns the - * CURRENT chunk, at the iteration's start] - * dbuf4 (this file): barrier -> prefetch(next -> temp) -> MMA(cur) -> - * store(temp -> next slice) [store owns the NEXT chunk, at the - * iteration's end -- the mirror image] - * - * The group wsum/wsc ping-pong is inverted the same way: dbuf4 stores the - * next group's values (prefetched during the crossing chunk) at the TAIL of - * that chunk, instead of dbuf2's HEAD-of-new-group placement. - * - * Selected at dispatch via - * ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4_txkgs<32|64> - * (QuantizedLinear.cpp), additive to the existing tsweep_t... (dbuf2) - * namespace. - * - * KHR Cooperative Matrix variant of the dynamically-quantized-activation - * linear tiled shader (WEIGHT_NBITS=4): - * 4 -> linear_dq8ca_q4gsw_coopmat INT4 group-symmetric weight - * - * Performs: out[M,N] = dequant(int8_act) * dequant(int_w) (+ bias) - * via coopmat x coopmat -> coopmat on the matrix unit. - * - * Hard preconditions: - * M % WG_TILE_M == 0, N % WG_TILE_N == 0, K % WG_TILE_K == 0, - * INT4: group_size % WG_TILE_K == 0, - * device exposes coopmatx-> at 16x16x16. - */ - -#version 450 core - -#extension GL_KHR_cooperative_matrix : require -#extension GL_KHR_memory_scope_semantics : require -#extension GL_KHR_shader_subgroup_basic : enable -#extension GL_EXT_shader_explicit_arithmetic_types : require -#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require -#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require -#extension GL_EXT_control_flow_attributes : enable - -#define PRECISION ${PRECISION} - -$if WEIGHT_NBITS == 4: - #define WEIGHT_INT4 - -// INTERVENTION G: when the A staging thread map exactly covers the workgroup -// (A_ACTIVE_THREADS == WG_SIZE) the `a_active` guard is statically always true, -// but the driver compiler does not fold it -- gl_LocalInvocationID.x's bound -// comes from a spec constant, so the comparison survives into the hot loop as a -// real branch. Set A_MAP_FULL only for tiles where the equality has been -// checked arithmetically; the yaml records the arithmetic per variant. -$if A_MAP_FULL: - #define A_ALWAYS_ACTIVE - -$if HAS_BIAS: - #define HAS_BIAS - -$if WEIGHT_STORAGE == "buffer": - #define WEIGHT_BUFFER - -$if IO_STORAGE == "texture3d": - #define IO_TEXTURE - -layout(std430) buffer; - -#include "common.glslh" - -// Bindings — match add_linear_dqa_qw_node arg order: -// output(0), fp_input(1), packed_int8_input(2), int_input_sums(3 - unused), -// input_scales(4), input_zps(5), packed_weight(6), weight_sums(7), -// weight_scales(8), bias(9). -${layout_declare_tensor(B, "w", "t_output", "half", IO_STORAGE, is_scalar_array=True)} -// t_input is unread here -- the activations arrive already quantized in -// t_packed_int8_input -- but stays declared so the binding layout matches the -// dispatch site. It tracks IO_STORAGE so the two IO tensors stay consistent. -${layout_declare_tensor(B, "r", "t_input", "half", IO_STORAGE, is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_packed_int8_input", "int", "buffer", is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_int8_input_sums", "int", "buffer", is_scalar_array=True)} -${layout_declare_tensor(B, "r", "t_int8_input_scales", "half", "texture3d")} -${layout_declare_tensor(B, "r", "t_int8_input_zps", "int8", "texture3d")} -${layout_declare_tensor(B, "r", "t_packed_weight", "int", WEIGHT_STORAGE, is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_weight_sums", "int", "buffer", is_scalar_array=True)} -${layout_declare_tensor(B, "r", "t_weight_scales", "half", "buffer", is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_bias", "half", "buffer", is_scalar_array=True)} - -${layout_declare_ubo(B, "ivec4", "output_sizes")} -${layout_declare_ubo(B, "ivec4", "input_sizes")} - -layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; - -${layout_declare_spec_const(C, "int", "apply_bias", "0")} -// INT4 only; inert (0) for INT8 so the dispatcher's spec list lines up. -${layout_declare_spec_const(C, "int", "K4_per_group", "0")} -${layout_declare_spec_const(C, "int", "num_groups_arg", "0")} -${layout_declare_spec_const(C, "int", "out_N_arg", "0")} - -// Tile geometry -const uint MMA_M = ${MMA_M}; -const uint MMA_N = ${MMA_N}; -const uint MMA_K = ${MMA_K}; - -const uint WG_TILE_M = ${WG_TILE_M}; -const uint WG_TILE_N = ${WG_TILE_N}; -const uint WG_TILE_K = ${WG_TILE_K}; - -const uint SG_GRID_X = ${SG_GRID_X}; -const uint SG_GRID_Y = ${SG_GRID_Y}; -const uint SUBGROUP_SIZE = ${SUBGROUP_SIZE}; -const uint NUM_SUBGROUPS = SG_GRID_X * SG_GRID_Y; -const uint WG_SIZE = NUM_SUBGROUPS * SUBGROUP_SIZE; - -const uint SG_TILE_M = WG_TILE_M / SG_GRID_Y; -const uint SG_TILE_N = WG_TILE_N / SG_GRID_X; -const uint MMAS_PER_SG_M = SG_TILE_M / MMA_M; -const uint MMAS_PER_SG_N = SG_TILE_N / MMA_N; - -const uint A_SLAB_INT8 = WG_TILE_M * MMA_K; -const uint B_USEFUL_U32 = MMA_K / 4u; -// No skew + coalesced write, ported from the production dbuf4.glsl fix -// (dq8ca-dequant-unpack-ablation Addenda 6/9): stride=4 beats stride=5 on -// this hardware, and writing via the inverted contiguous-index mapping -// (see the F-hoist block below) beats the natural strided-scatter mapping. -const uint B_STRIDE_U32 = B_USEFUL_U32; -const uint B_SLAB_U32 = WG_TILE_N * B_STRIDE_U32; -const uint NUM_K_SLABS = WG_TILE_K / MMA_K; - -const uint A_SLAB_U32 = A_SLAB_INT8 / 4u; -const uint A_STRIDE_U32 = MMA_K / 4u; - -// One ping-pong slice covers all K-slabs of one chunk. -const uint ASH_SLICE_U32 = NUM_K_SLABS * A_SLAB_U32; -const uint BSH_SLICE_U32 = NUM_K_SLABS * B_SLAB_U32; - -// Double-buffered MMA operand staging. -shared uint Ash_int8[2u * ASH_SLICE_U32]; -shared uint Bsh_int8[2u * BSH_SLICE_U32]; - -// Per-WG-tile-row activation params (loaded ONCE at WG start; constant -// across groups). -shared int izp_sh[WG_TILE_M]; // int32 (cast from int8 source) for broadcast -shared float ifs_sh[WG_TILE_M]; // float32 (cast from fp16 source) for broadcast - -// Per-(group, output-channel) weight params, ping-ponged by group parity. -// (For per-channel INT8 only slice 0 is ever used.) -shared float wsc_sh[2u * WG_TILE_N]; -// SUM_g wsc[g][n]*wsum[g][n] per output channel -- weight-side only, so it is -// accumulated once in the prologue. Replaces dbuf4's ping-ponged wsum_sh -// (which was 2*WG_TILE_N ints), so this is a net LDS saving. -shared float wcorr_sh[WG_TILE_N]; - -#ifdef HAS_BIAS -shared float bias_sh[WG_TILE_N]; -#endif - -#ifdef IO_TEXTURE -// Result staging for the imageStore epilogue, mirroring the fp16 kernel: -// SG_GRID_Y bands of MMA_M rows, each WG_TILE_N wide, row-major. A full -// WG_TILE_M x WG_TILE_N buffer would cost SG_GRID_Y/MMAS_PER_SG_M x more LDS -// and wreck occupancy. float16_t-typed because coopMatStore needs it. -const uint CSH_ROWS = SG_GRID_Y * MMA_M; -shared float16_t Csh_out[CSH_ROWS * WG_TILE_N]; -#endif - -// Running fp32 accumulator (across all groups). -coopmat - result[MMAS_PER_SG_M][MMAS_PER_SG_N]; - -// Per-group int32 MMA accumulator. -coopmat - accum_int32[MMAS_PER_SG_M][MMAS_PER_SG_N]; - - -// Byte-parallel int4 -> int8 widening. -// -// The four nibbles this shader needs from one packed uint are ALREADY one per -// byte (bits 3:0 of each byte for parity 0, bits 7:4 for parity 1), so all four -// can be widened at once instead of with a per-nibble -// shift/mask/bias-subtract/mask chain. -// -// For v in [0,15] the biased value is v-8. `v ^ 8` is exactly the 4-bit -// two's-complement pattern of v-8, because -8 == +8 (mod 16): -// v=0 -> 0x8 -> -8 v=7 -> 0xF -> -1 -// v=8 -> 0x0 -> 0 v=15 -> 0x7 -> +7 -// so the only remaining work is sign-extending bit 3 into bits 7:4 per byte. -// `sgn * 0x1E` does that with no cross-byte carry: 0x08 * 0x1E == 0xF0 exactly, -// and sgn is at most 0x08080808 so the product is at most 0xF0F0F0F0. -// -// A naive `nib - 0x08080808` would NOT work -- it borrows across byte lanes -// whenever a nibble is < 8. Shifts must be on uint, not int, so the >> is -// logical rather than arithmetic. -// -// ~5 ops per 4 weights vs ~22 for the per-nibble chain; bit-identical output. -uint widen_nibbles(const uint w, const uint parity) { - const uint nib = (parity == 0u) ? (w & 0x0F0F0F0Fu) : ((w >> 4u) & 0x0F0F0F0Fu); - const uint p = nib ^ 0x08080808u; - const uint sgn = p & 0x08080808u; - return p | (sgn * 0x1Eu); -} - -void main() { - const uvec2 tileID = uvec2(gl_WorkGroupID.xy); - const uvec2 warpInTile = uvec2( - gl_SubgroupID % SG_GRID_X, - gl_SubgroupID / SG_GRID_X); - - const uint K = uint(input_sizes.x); - const uint N = uint(output_sizes.x); - const uint N4 = (N + 3u) / 4u; - const uint nblocks_x_A = (K + 3u) >> 2u; - -#ifdef WEIGHT_INT4 - const uint num_groups = uint(num_groups_arg); - const uint CHUNKS_PER_GROUP = uint(K4_per_group) * 4u / WG_TILE_K; -#else - // Per-channel: a single quant "group" spanning all of K. The nested - // groups x chunks loop below collapses to a flat chunk loop, the wsum/wsc - // ping-pong never crosses a boundary, and the epilog runs exactly once. - const uint num_groups = 1u; - const uint CHUNKS_PER_GROUP = uint(num_groups_arg); -#endif - const uint num_chunks = num_groups * CHUNKS_PER_GROUP; - - const uint tile_m_start = WG_TILE_M * tileID.y; - const uint tile_n_start = WG_TILE_N * tileID.x; - - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - result[i][j] = coopmat(0.0); - accum_int32[i][j] = coopmat(0); - } - } - - // --- A staging thread map: one (m4, k4) ivec4 block per active thread --- - const uint K_BLOCKS_PER_CHUNK = WG_TILE_K >> 2u; - const uint A_ACTIVE_THREADS = (WG_TILE_M >> 2u) * K_BLOCKS_PER_CHUNK; - const uint a_m_block = gl_LocalInvocationID.x / K_BLOCKS_PER_CHUNK; - const uint a_k_block = gl_LocalInvocationID.x % K_BLOCKS_PER_CHUNK; -#ifdef A_ALWAYS_ACTIVE - // A_ACTIVE_THREADS == WG_SIZE for this variant's tile, so every thread stages - // A and the guard is unconditionally true. - const bool a_active = true; -#else - const bool a_active = gl_LocalInvocationID.x < A_ACTIVE_THREADS; -#endif - -#ifdef WEIGHT_INT4 - // --- B staging thread map: (block, col) slots; each slot extracts one - // ColumnMajor LDS uint (4 K-contiguous sign-extended int8) --- - const uint B_TOTAL_SLOTS = K_BLOCKS_PER_CHUNK * WG_TILE_N; - const uint B_SLOTS_PER_THREAD = B_TOTAL_SLOTS / WG_SIZE; - const uint N8_PER_TILE = WG_TILE_N >> 3u; -#else - // --- B staging thread map: one (k4, n4) ivec4 block per active thread --- - const uint B_FETCH_SLOTS = K_BLOCKS_PER_CHUNK * (WG_TILE_N >> 2u); - const uint N4_PER_TILE = WG_TILE_N >> 2u; - const uint b_k4_in_chunk = gl_LocalInvocationID.x / N4_PER_TILE; - const uint b_n_uint_col = gl_LocalInvocationID.x % N4_PER_TILE; - const bool b_active = gl_LocalInvocationID.x < B_FETCH_SLOTS; -#endif - - // ===== INTERVENTION F: hoist loop-invariant staging index math ===== - // Every term below depends only on gl_LocalInvocationID.x, si, tile_n_start - // and tile_m_start -- none of them on `chunk` or `group_i`. In the baseline - // this same 10-line decomposition is recomputed at THREE sites (prologue - // store, in-loop prefetch, in-loop store) on every chunk iteration. The - // ablation ladder attributes -16.8% to the B staging block and shows that - // most of it is index arithmetic, not memory traffic: deleting B staging - // removed 45 of the loop's 96 address instructions and 56 of its 64 - // compare/select instructions, but only 24 memory ops. - // - // Only the k4 texel row varies with the chunk, and it varies by a constant - // stride, so the loop carries an increment instead of a recomputation. -#ifdef WEIGHT_INT4 - uint b_lds_off[B_SLOTS_PER_THREAD]; // LDS store offset within a slice - uint b_comp[B_SLOTS_PER_THREAD]; // which ivec4 component feeds this slot - uint b_par[B_SLOTS_PER_THREAD]; // nibble parity for this slot - uint b_n8blk[B_SLOTS_PER_THREAD]; // global texel column (N/8 blocks) - uint b_k4off[B_SLOTS_PER_THREAD]; // k4 offset of this slot within a chunk - // Coalesced-write inversion (ported from production dbuf4.glsl's - // bcoal_index, dq8ca-dequant-unpack-ablation Addendum 9): start from the - // CONTIGUOUS per-thread LDS index `a` this thread will write to, and derive - // which global-fetch element it needs -- instead of deriving the LDS - // address from the natural fetch grouping (which is what produced the - // B_STRIDE_U32-strided, non-coalesced writes this fix replaces). - [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { - const uint a = gl_LocalInvocationID.x + si * WG_SIZE; - const uint slab_idx = a / B_SLAB_U32; - const uint local_a = a % B_SLAB_U32; - const uint n_col = local_a / B_STRIDE_U32; - const uint k4_in_slab = local_a % B_STRIDE_U32; - const uint k4_in_chunk = slab_idx * (MMA_K >> 2u) + k4_in_slab; - const uint n8_in_tile = n_col >> 3u; - const uint rem = n_col & 7u; - b_lds_off[si] = a; - b_comp[si] = rem & 3u; - b_par[si] = rem >> 2u; - b_n8blk[si] = (tile_n_start >> 3u) + n8_in_tile; - b_k4off[si] = k4_in_chunk; - } -#endif - // A staging: same argument. base_row/slab/k_uint are all invariant. - const uint a_lds_off0 = - (a_k_block / (MMA_K >> 2u)) * A_SLAB_U32 - + (a_m_block * 4u) * A_STRIDE_U32 - + (a_k_block % (MMA_K >> 2u)); - const uint a_glb_row = ((tile_m_start >> 2u) + a_m_block) * nblocks_x_A; - - // Prefetch temp registers. - ivec4 temp_A; -#ifdef WEIGHT_INT4 - ivec4 temp_B[B_SLOTS_PER_THREAD]; - float temp_wsc; -#else - ivec4 temp_B; -#endif - - // ========================================================= - // PROLOGUE - // ========================================================= - if (gl_LocalInvocationID.x < (WG_TILE_M >> 2u)) { - const uint m4 = (tile_m_start >> 2u) + gl_LocalInvocationID.x; - const vec4 sc = vec4(texelFetch(t_int8_input_scales, ivec3(m4, 0, 0), 0)); - const ivec4 zp = texelFetch(t_int8_input_zps, ivec3(m4, 0, 0), 0); - const uint base = gl_LocalInvocationID.x * 4u; - ifs_sh[base + 0u] = sc.x; ifs_sh[base + 1u] = sc.y; - ifs_sh[base + 2u] = sc.z; ifs_sh[base + 3u] = sc.w; - izp_sh[base + 0u] = zp.x; izp_sh[base + 1u] = zp.y; - izp_sh[base + 2u] = zp.z; izp_sh[base + 3u] = zp.w; - } - // Group 0 weight scales -> slice 0, and the hoisted weight-side correction - // SUM_g wsc[g][n]*wsum[g][n] accumulated across ALL groups. The loop is - // prologue-only (the prologue is ~1.4% of the dynamic instruction stream), - // and it replaces per-group wsum work inside the loop body. - if (gl_LocalInvocationID.x < WG_TILE_N) { - const uint n_idx = tile_n_start + gl_LocalInvocationID.x; - f16vec4 sv0 = t_weight_scales[n_idx >> 2u]; - wsc_sh[gl_LocalInvocationID.x] = float(sv0[n_idx & 3u]); - - float corr = 0.0; - for (uint g = 0; g < num_groups; ++g) { - f16vec4 sv = t_weight_scales[g * N4 + (n_idx >> 2u)]; - corr += float(sv[n_idx & 3u]) * float(t_weight_sums[g * N + n_idx]); - } - wcorr_sh[gl_LocalInvocationID.x] = corr; - } - memoryBarrierShared(); - barrier(); - - // NOTE: dbuf4 builds izp_bcast/ifs_bcast here and keeps them live across the - // whole group loop. This variant needs them only AFTER the loop, so they are - // loaded there instead -- that is the register-pressure saving. - - // dbuf4: prefetch chunk 0 into temp registers, THEN store to slice 0 (no - // barrier here -- the main loop's first iteration barriers before - // reading slice 0). - if (a_active) { - const uint m4_global = (tile_m_start >> 2u) + a_m_block; - temp_A = t_packed_int8_input[m4_global * nblocks_x_A + a_k_block]; - } -#ifdef WEIGHT_INT4 - [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { -#ifdef WEIGHT_BUFFER - temp_B[si] = t_packed_weight[(b_n8blk[si] * nblocks_x_A) + b_k4off[si]]; -#else - temp_B[si] = texelFetch(t_packed_weight, ivec2(b_k4off[si], b_n8blk[si]), 0); -#endif - } -#else - if (b_active) { - const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; -#ifdef WEIGHT_BUFFER - temp_B = t_packed_weight[(b_k4_in_chunk * N4) + block_x_w]; -#else - temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, b_k4_in_chunk), 0); -#endif - } -#endif - { - // store chunk 0 -> slice 0 -#ifndef A_ALWAYS_ACTIVE - if (a_active) -#endif - { - [[unroll]] for (uint m4i = 0; m4i < 4u; ++m4i) { - Ash_int8[a_lds_off0 + m4i * A_STRIDE_U32] = uint(temp_A[m4i]); - } - } -#ifdef WEIGHT_INT4 - [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { - Bsh_int8[b_lds_off[si]] = - widen_nibbles(uint(temp_B[si][b_comp[si]]), b_par[si]); - } -#else - if (b_active) { - const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); - const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); - const uint n_col_base = b_n_uint_col * 4u; - [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { - Bsh_int8[slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab] = - uint(temp_B[n_in_blk]); - } - } -#endif - } - - // ========================================================= - // MAIN LOOP (dbuf4) — nested groups x chunks (kept nested; flattening it - // with a conditional coopmat epilog crashes the Xclipse PAL compiler at - // large spec-resolved trip counts). One barrier per chunk. Chunk - // iteration (global index `chunk`): - // 1. barrier — A/B slice (chunk%2) fully written; on the first chunk - // of group g, wsum/wsc slice (g%2) is too. - // 2. prefetch — chunk+1 (A blocks, B blocks) into temp; when chunk+1 - // starts a new group, also its wsum/wsc element. Skipped - // entirely on the final chunk. - // 3. int8 MMA — on slice (chunk%2) into accum_int32. - // 4. store — temp -> A/B slice ((chunk+1)%2), unpacking the weight; - // on a group boundary, wsum/wsc -> slice ((g+1)%2). - // The group epilog runs unconditionally at the tail of each group. - // ========================================================= - uint chunk = 0; - for (uint group_i = 0; group_i < num_groups; ++group_i) { - for (uint inner = 0; inner < CHUNKS_PER_GROUP; ++inner, ++chunk) { - const bool has_next = chunk + 1u < num_chunks; - const bool group_crossing = has_next && (inner + 1u == CHUNKS_PER_GROUP); - const uint cur_a = (chunk % 2u) * ASH_SLICE_U32; - const uint cur_b = (chunk % 2u) * BSH_SLICE_U32; - const uint nxt_a = ((chunk + 1u) % 2u) * ASH_SLICE_U32; - const uint nxt_b = ((chunk + 1u) % 2u) * BSH_SLICE_U32; - - // coopmat-lds-fence: barrier() alone does NOT order shared stores against a - // subsequent coopMatLoad on the M51 Xclipse/AMD-PAL driver -- symptom is one - // stale MMA_M-row band of A, all columns, ~2.5% of runs, no crash (observed - // 2026-09-02 in sdpa_compute_out_coopmat.glsl). Measured cost of the fence: - // none (see this change's results). See memory - // `coopmat-lds-needs-explicit-memorybarriershared`. - memoryBarrierShared(); - barrier(); - - // --- 2. prefetch chunk+1 -> temp --- - if (has_next) { - const uint chunkK_nxt = (chunk + 1u) * WG_TILE_K; -#ifndef A_ALWAYS_ACTIVE - if (a_active) -#endif - { - temp_A = t_packed_int8_input[a_glb_row + (chunkK_nxt >> 2u) + a_k_block]; - } -#ifdef WEIGHT_INT4 - [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { - const uint k4_blk = (chunkK_nxt >> 2u) + b_k4off[si]; -#ifdef WEIGHT_BUFFER - temp_B[si] = t_packed_weight[(b_n8blk[si] * nblocks_x_A) + k4_blk]; -#else - temp_B[si] = texelFetch(t_packed_weight, ivec2(k4_blk, b_n8blk[si]), 0); -#endif - } - if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { - const uint n_idx = tile_n_start + gl_LocalInvocationID.x; - f16vec4 sv = t_weight_scales[(group_i + 1u) * N4 + (n_idx >> 2u)]; - temp_wsc = float(sv[n_idx & 3u]); - } -#else - if (b_active) { - const uint block_y_w = (chunkK_nxt >> 2u) + b_k4_in_chunk; - const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; -#ifdef WEIGHT_BUFFER - temp_B = t_packed_weight[(block_y_w * N4) + block_x_w]; -#else - temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, block_y_w), 0); -#endif - } -#endif - } - - // --- 3. int8 MMA on the cur slice --- - [[unroll]] for (uint k = 0; k < NUM_K_SLABS; ++k) { - const uint slab_a_base_u32 = cur_a + k * A_SLAB_U32; - const uint slab_b_base_u32 = cur_b + k * B_SLAB_U32; - - coopmat matA[MMAS_PER_SG_M]; - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - const uint row_a = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); - coopMatLoad( - matA[i], Ash_int8, - slab_a_base_u32 + row_a * A_STRIDE_U32, - A_STRIDE_U32, - gl_CooperativeMatrixLayoutRowMajor); - } - - coopmat matB; - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint col_b = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopMatLoad( - matB, Bsh_int8, - slab_b_base_u32 + col_b * B_STRIDE_U32, - B_STRIDE_U32, - gl_CooperativeMatrixLayoutColumnMajor); - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - accum_int32[i][j] = coopMatMulAdd(matA[i], matB, accum_int32[i][j]); - } - } - } - - // --- 4. store temp (chunk+1) -> nxt slice --- - if (has_next) { -#ifndef A_ALWAYS_ACTIVE - if (a_active) -#endif - { - [[unroll]] for (uint m4i = 0; m4i < 4u; ++m4i) { - Ash_int8[nxt_a + a_lds_off0 + m4i * A_STRIDE_U32] = uint(temp_A[m4i]); - } - } -#ifdef WEIGHT_INT4 - [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { - Bsh_int8[nxt_b + b_lds_off[si]] = - widen_nibbles(uint(temp_B[si][b_comp[si]]), b_par[si]); - } - if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { - const uint wbase_nxt = ((group_i + 1u) % 2u) * WG_TILE_N; - wsc_sh[wbase_nxt + gl_LocalInvocationID.x] = temp_wsc; - } -#else - if (b_active) { - const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); - const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); - const uint n_col_base = b_n_uint_col * 4u; - [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { - Bsh_int8[nxt_b + slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab] = - uint(temp_B[n_in_blk]); - } - } -#endif - } - } // chunks - - // --- Group epilog: scale-only accumulate, reset accum --- - // Just result += float(acc) * wsc. The zero-point subtract and the ifs - // multiply are hoisted out of the group loop (applied once below). - { - const uint wbase = (group_i % 2u) * WG_TILE_N; - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - - coopmat wsc_bcast; - coopMatLoad( - wsc_bcast, wsc_sh, - wbase + local_n_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutRowMajor); - - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - result[i][j] += - coopmat( - accum_int32[i][j]) * wsc_bcast; - accum_int32[i][j] = coopmat(0); - } - } - } - } // groups - - // --- Hoisted correction, applied ONCE: --------------------------------- - // result = ifs * ( result - izp * SUM_g wsc_g*wsum_g ) - // izp/ifs are loaded here rather than before the group loop so they are not - // live across it. - { - coopmat - izpf_bcast[MMAS_PER_SG_M]; - coopmat - ifs_bcast[MMAS_PER_SG_M]; - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - const uint local_m_base = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); - coopmat izp_i; - coopMatLoad( - izp_i, izp_sh, local_m_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutColumnMajor); - izpf_bcast[i] = - coopmat(izp_i); - coopMatLoad( - ifs_bcast[i], ifs_sh, local_m_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutColumnMajor); - } - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopmat wcorr_bcast; - coopMatLoad( - wcorr_bcast, wcorr_sh, local_n_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutRowMajor); - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - result[i][j] = ifs_bcast[i] * (result[i][j] - izpf_bcast[i] * wcorr_bcast); - } - } - } - - // --- Bias (optional) --- -#ifdef HAS_BIAS - if (apply_bias > 0) { - for (uint t = gl_LocalInvocationID.x; t < WG_TILE_N; t += WG_SIZE) { - bias_sh[t] = float(t_bias[tile_n_start + t]); - } - memoryBarrierShared(); - barrier(); - } -#endif - - // --- Store result tile --- - // N for the store address math MUST come from the spec constant, not the - // sizes UBO (see out_N_arg above). -#ifdef IO_TEXTURE - // Epilogue iteration i drains accumulator row-block i from EVERY subgroup - // into Csh_out at once, so the SG_GRID_Y bands it holds are disjoint global - // row ranges; the whole workgroup then imageStores them. lr / MMA_M is the - // writing subgroup's warpInTile.y, so the global row reproduces the buffer - // path's gi exactly. - // - // PORTABILITY NOTE: the barrier() in the loop body keeps this loop rolled - // despite [[unroll]], so result[i][j] IS dynamically indexed. Coopmat arrays - // are opaque per-lane storage and dynamic indexing is exactly the construct - // the Xclipse/AMD-PAL compiler has broken before -- check this first if the - // texture variants miscompile on M51. - const uint CSH_TEXELS_PER_ROW = WG_TILE_N / 4u; - const uint CSH_TEXELS = CSH_ROWS * CSH_TEXELS_PER_ROW; - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - // Guards Csh_out against the previous iteration's readers. Inert on i == 0 - // but must stay unconditional to remain workgroup-uniform. - // coopmat-lds-fence: barrier() alone does NOT order shared stores against a - // subsequent coopMatLoad on the M51 Xclipse/AMD-PAL driver -- symptom is one - // stale MMA_M-row band of A, all columns, ~2.5% of runs, no crash (observed - // 2026-09-02 in sdpa_compute_out_coopmat.glsl). Measured cost of the fence: - // none (see this change's results). See memory - // `coopmat-lds-needs-explicit-memorybarriershared`. - memoryBarrierShared(); - barrier(); - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { -#ifdef HAS_BIAS - if (apply_bias > 0) { - const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopmat bias_tile; - coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); - result[i][j] += bias_tile; - } -#endif - coopmat out_tile = - coopmat(result[i][j]); - coopMatStore( - out_tile, Csh_out, - warpInTile.y * MMA_M * WG_TILE_N + - MMA_N * (MMAS_PER_SG_N * warpInTile.x + j), - WG_TILE_N, - gl_CooperativeMatrixLayoutRowMajor); - } - memoryBarrierShared(); - barrier(); - - for (uint t = gl_LocalInvocationID.x; t < CSH_TEXELS; t += WG_SIZE) { - const uint lr = t / CSH_TEXELS_PER_ROW; - const uint lc4 = t % CSH_TEXELS_PER_ROW; - const uint m = - tile_m_start + (lr / MMA_M) * SG_TILE_M + i * MMA_M + (lr % MMA_M); - const uint base = lr * WG_TILE_N + lc4 * 4u; - imageStore( - t_output, - ivec3(tile_n_start / 4u + lc4, m, 0), - vec4( - float(Csh_out[base]), - float(Csh_out[base + 1u]), - float(Csh_out[base + 2u]), - float(Csh_out[base + 3u]))); - } - } -#else - const uint N_out = uint(out_N_arg); - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint gi = tile_m_start + MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); - const uint gj = tile_n_start + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - -#ifdef HAS_BIAS - if (apply_bias > 0) { - const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopmat bias_tile; - coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); - result[i][j] += bias_tile; - } -#endif - - coopmat out_tile = - coopmat(result[i][j]); - coopMatStore( - out_tile, t_output, - gi * N_out + gj, N_out, - gl_CooperativeMatrixLayoutRowMajor); - } - } -#endif // IO_TEXTURE -} diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpg.yaml b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpg.yaml deleted file mode 100644 index 444c8347cde..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpg.yaml +++ /dev/null @@ -1,80 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -# "zp-hoisted" variants of the int8 dq8ca_q4gsw coopmat kernel: identical to -# tsweep_dbuf4 except the activation zero-point correction and the per-row -# activation scale are applied once after the group loop instead of once per -# quantization group. No new binding, no export-format change -- the weight-side -# correction sum is derived in the prologue from tensors already bound. -# -# Selected via -# ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4zpg_txkgs. -# NOT the shipped default until it passes repeated correctness runs. - -linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpg: - parameter_names_with_default_values: - PRECISION: highp - HAS_BIAS: false - IO_STORAGE: buffer - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - A_MAP_FULL: false - MMA_M: 16 - MMA_N: 16 - MMA_K: 16 - WG_TILE_M: 64 - WG_TILE_N: 32 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - shader_variants: - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpg_t128x128k64g81s64_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 8 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 64 - # A_ACTIVE_THREADS = (128>>2)*(64>>2) = 512 == WG_SIZE = 8*1*64 = 512 - A_MAP_FULL: true - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpg_t128x128k64g81s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 8 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 64 - A_MAP_FULL: true - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpg_t128x64k32g42s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - # A_ACTIVE_THREADS = (128>>2)*(32>>2) = 256 == WG_SIZE = 4*2*32 = 256 - A_MAP_FULL: true - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpg_t128x64k32g42s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - # small step up - # A_ACTIVE_THREADS = (128>>2)*(32>>2) = 256 == WG_SIZE = 4*2*32 = 256 - A_MAP_FULL: true diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbv4.glsl b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbv4.glsl deleted file mode 100644 index c74d6558055..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbv4.glsl +++ /dev/null @@ -1,792 +0,0 @@ -/* - * Copyright (c) Meta Platforms, Inc. and affiliates. - * All rights reserved. - * - * This source code is licensed under the BSD-style license found in the - * LICENSE file in the root directory of this source tree. - */ - -/* - * "zpg" + "-tr" combination: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpg.glsl - * with its per-thread scalar A-staging replaced by - * linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr.glsl's coopMat-mediated A - * staging, PLUS "bv4": shared B retyped to uvec4 so each (K-slab, column) - * column-slab is ONE 16-byte element written by ONE b128 store instead of 4 - * scalar ds_write_b32 spread across 4 threads. A staging, the quantization - * math and every correctness safeguard are dbuf4zpgtr's, unchanged. This is - * an ADDITIVE combination, not a redesign: every non-A-staging block below - * (B staging: coalesced write, no skew; zp-hoist: izp/ifs applied once after - * the group loop via wcorr_sh; byte-parallel nibble widening; static - * A_ALWAYS_ACTIVE branch elision -- N/A here, see below; group epilog; - * bias/store epilogue) is byte-identical to dbuf4zpg's. Only the A-staging - * block (prologue load+store, main-loop prefetch+store) is dbuf4tr's, - * verbatim. - * - * dbuf4zpg's per-thread A staging used an `a_active` guard (statically always - * true when A_ACTIVE_THREADS == WG_SIZE, via the A_MAP_FULL-gated - * A_ALWAYS_ACTIVE macro). dbuf4tr's per-SUBGROUP tile map has no equivalent - * concept -- every subgroup participates via a `t < NUM_A_TILES` guard that - * depends only on gl_SubgroupID, not gl_LocalInvocationID.x -- so - * A_MAP_FULL/A_ALWAYS_ACTIVE is dropped entirely in this file; it would be - * dead code for the new A-staging block. - * - * Rationale for combining this way (not the reverse) and why this is worth - * building at all: see this change's design.md D0-D3. In short -- the only - * existing measurement of dbuf4tr's A-staging technique (28.72-30.51%, - * dq8ca-arch-redesign) was taken against dbuf4tr's own pre-zpg baseline - * (old B skew, no byte-parallel widening, no branch elision) -- a materially - * weaker shader than the 46.49-46.50% dbuf4zpg this file now combines it - * with. This file exists to answer whether that combination performs - * differently now that register pressure is already reduced. - * - * A staging (the actual delta from dbuf4zpg): - * dbuf4zpg: per-thread (m4, k4) ivec4 fetch, hoisted a_lds_off0/a_glb_row; - * only A_ACTIVE_THREADS invocations participate, each scattering - * 4 rows into Ash_int8 with 4 scalar stores. - * this file: per-SUBGROUP MMA_M x MMA_K tile fetch via coopMatLoad straight - * from a ROW-MAJOR (kPackedInt8_4W) int8 activation buffer, then - * coopMatStore into the same Ash_int8 slot -- dbuf4tr's mapping, - * unmodified (not re-derived; see design.md D3). - * - * t_packed_int8_input is therefore bound the same way dbuf4tr binds it: a - * SCALAR int8_t array in the kPackedInt8_4W layout (plain row-major int8, - * row stride K), produced by quantize_and_pack_4w_with_group_sums.glsl. - * QuantizedLinear.cpp's dq8ca_variant_wants_rowmajor_a() must recognize this - * file's variant token (tsweep_dbuf4zpgtr_t...) the same way it already - * recognizes tsweep_dbuf4tr_t/trm_t/trd_t, so graph-build time (packer - * selection) and dispatch time (kernel selection) cannot disagree. - * - * B CANNOT be coopmat-staged (int4 nibble unpack; a coopmat's per-lane layout - * is opaque to hand-assembly from unpacked registers) -- unchanged from both - * parent files. B staging below is dbuf4zpg's byte-parallel, coalesced, - * no-skew version, untouched. - * - * The loop structure is dbuf4's (both parents share it), unchanged: - * prologue: prefetch chunk 0 -> temp, store to slice 0 (no barrier) - * per iter: barrier -> prefetch(next) -> MMA(cur) -> store(next) - * kept nested (groups x chunks) with an unconditional group epilog -- - * flattening it crashes the Xclipse PAL compiler at large spec-resolved trip - * counts (see dbuf2's own header). - * - * Selected via - * ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4zpgtr_txkgs<32|64> - * (QuantizedLinear.cpp), additive to the tsweep_dbuf4zpg_t..., tsweep_dbuf4tr_t... - * and tsweep_t... namespaces. NOT the default -- unvalidated until it passes - * repeated test_llama_microbench --correctness-only runs (see - * dq8ca_coopmat_variant()'s comment on why a single pass is not proof). - * - * Performs: out[M,N] = dequant(int8_act) * dequant(int_w) (+ bias) - * via coopmat x coopmat -> coopmat on the matrix unit. - * - * Hard preconditions (dbuf4zpg's, plus dbuf4tr's row-major/alignment ones): - * M % WG_TILE_M == 0, N % WG_TILE_N == 0, K % WG_TILE_K == 0, - * group_size % WG_TILE_K == 0, K % 4 == 0, - * WG_TILE_M % MMA_M == 0, WG_TILE_K % MMA_K == 0, - * t_packed_int8_input in kPackedInt8_4W (row-major) layout, - * device exposes coopmatx-> at 16x16x16. - */ - -#version 450 core - -#extension GL_KHR_cooperative_matrix : require -#extension GL_KHR_memory_scope_semantics : require -#extension GL_KHR_shader_subgroup_basic : enable -#extension GL_EXT_shader_explicit_arithmetic_types : require -#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require -// 8-bit SSBO access: A is bound as a scalar int8_t array so that the -// coopMatLoad below has a MATCHING component type (see dbuf4tr's header for -// why the type must match on this driver). -#extension GL_EXT_shader_8bit_storage : require -#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require -#extension GL_EXT_control_flow_attributes : enable - -#define PRECISION ${PRECISION} - -$if WEIGHT_NBITS == 4: - #define WEIGHT_INT4 - -$if HAS_BIAS: - #define HAS_BIAS - -$if WEIGHT_STORAGE == "buffer": - #define WEIGHT_BUFFER - -$if IO_STORAGE == "texture3d": - #define IO_TEXTURE - -layout(std430) buffer; - -#include "common.glslh" - -// Bindings — match add_linear_dqa_qw_node arg order: -// output(0), fp_input(1), packed_int8_input(2), int_input_sums(3 - unused), -// input_scales(4), input_zps(5), packed_weight(6), weight_sums(7), -// weight_scales(8), bias(9). -${layout_declare_tensor(B, "w", "t_output", "half", IO_STORAGE, is_scalar_array=True)} -// t_input is unread here -- the activations arrive already quantized in -// t_packed_int8_input -- but stays declared so the binding layout matches the -// dispatch site. It tracks IO_STORAGE so the two IO tensors stay consistent. -${layout_declare_tensor(B, "r", "t_input", "half", IO_STORAGE, is_scalar_array=False)} -// ROW-MAJOR (kPackedInt8_4W) packed activations, bound as a scalar int8_t -// array (row stride = K int8) -- dbuf4tr's binding, unchanged. The stock -// 4h4w layout dbuf4zpg uses is NOT row-major (component index selects a row, -// non-affine), so it cannot be addressed by any coopMatLoad. -${layout_declare_tensor(B, "r", "t_packed_int8_input", "int8", "buffer", is_scalar_array=True)} -${layout_declare_tensor(B, "r", "t_int8_input_sums", "int", "buffer", is_scalar_array=True)} -${layout_declare_tensor(B, "r", "t_int8_input_scales", "half", "texture3d")} -${layout_declare_tensor(B, "r", "t_int8_input_zps", "int8", "texture3d")} -${layout_declare_tensor(B, "r", "t_packed_weight", "int", WEIGHT_STORAGE, is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_weight_sums", "int", "buffer", is_scalar_array=True)} -${layout_declare_tensor(B, "r", "t_weight_scales", "half", "buffer", is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_bias", "half", "buffer", is_scalar_array=True)} - -${layout_declare_ubo(B, "ivec4", "output_sizes")} -${layout_declare_ubo(B, "ivec4", "input_sizes")} - -layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; - -${layout_declare_spec_const(C, "int", "apply_bias", "0")} -// INT4 only; inert (0) for INT8 so the dispatcher's spec list lines up. -${layout_declare_spec_const(C, "int", "K4_per_group", "0")} -${layout_declare_spec_const(C, "int", "num_groups_arg", "0")} -${layout_declare_spec_const(C, "int", "out_N_arg", "0")} - -// Tile geometry -const uint MMA_M = ${MMA_M}; -const uint MMA_N = ${MMA_N}; -const uint MMA_K = ${MMA_K}; - -const uint WG_TILE_M = ${WG_TILE_M}; -const uint WG_TILE_N = ${WG_TILE_N}; -const uint WG_TILE_K = ${WG_TILE_K}; - -const uint SG_GRID_X = ${SG_GRID_X}; -const uint SG_GRID_Y = ${SG_GRID_Y}; -const uint SUBGROUP_SIZE = ${SUBGROUP_SIZE}; -const uint NUM_SUBGROUPS = SG_GRID_X * SG_GRID_Y; -const uint WG_SIZE = NUM_SUBGROUPS * SUBGROUP_SIZE; - -const uint SG_TILE_M = WG_TILE_M / SG_GRID_Y; -const uint SG_TILE_N = WG_TILE_N / SG_GRID_X; -const uint MMAS_PER_SG_M = SG_TILE_M / MMA_M; -const uint MMAS_PER_SG_N = SG_TILE_N / MMA_N; - -const uint A_SLAB_INT8 = WG_TILE_M * MMA_K; -const uint B_USEFUL_U32 = MMA_K / 4u; -// No skew + coalesced write -- dbuf4zpg's B fix, unchanged (this file does -// not touch B staging at all). -const uint B_STRIDE_U32 = B_USEFUL_U32; -const uint B_SLAB_U32 = WG_TILE_N * B_STRIDE_U32; -const uint NUM_K_SLABS = WG_TILE_K / MMA_K; - -const uint A_SLAB_U32 = A_SLAB_INT8 / 4u; -const uint A_STRIDE_U32 = MMA_K / 4u; - -// One ping-pong slice covers all K-slabs of one chunk. -const uint ASH_SLICE_U32 = NUM_K_SLABS * A_SLAB_U32; -const uint BSH_SLICE_U32 = NUM_K_SLABS * B_SLAB_U32; - -// Double-buffered MMA operand staging. -shared uint Ash_int8[2u * ASH_SLICE_U32]; -// ===== bv4: shared B is uvec4-typed, ONE element per (K-slab, column) ===== -// dbuf4zpgtr stores B as `uint`: 4 uints per column-slab, written as 4 separate -// scalar ds_write_b32 spread over 4 threads. One column-slab is exactly -// MMA_K int8 = 16 bytes = one uvec4, and B_STRIDE_U32 is already 4 because zpg -// removed the anti-bank-conflict skew -- so those 4 uints are contiguous AND -// 16B-aligned. That alignment is what made a b128 store impossible on the old -// skewed layout (75% of columns were unaligned) and possible now. -// Each owning thread assembles all 4 widened uints and issues ONE store. -const uint BSH_SLICE_V4 = NUM_K_SLABS * WG_TILE_N; -shared uvec4 Bsh_v4[2u * BSH_SLICE_V4]; - -// Per-WG-tile-row activation params (loaded ONCE at WG start; constant -// across groups). -shared int izp_sh[WG_TILE_M]; // int32 (cast from int8 source) for broadcast -shared float ifs_sh[WG_TILE_M]; // float32 (cast from fp16 source) for broadcast - -// Per-(group, output-channel) weight params, ping-ponged by group parity. -// (For per-channel INT8 only slice 0 is ever used.) -shared float wsc_sh[2u * WG_TILE_N]; -// SUM_g wsc[g][n]*wsum[g][n] per output channel -- weight-side only, so it is -// accumulated once in the prologue. dbuf4zpg's zp-hoist, unchanged. -shared float wcorr_sh[WG_TILE_N]; - -#ifdef HAS_BIAS -shared float bias_sh[WG_TILE_N]; -#endif - -#ifdef IO_TEXTURE -// Result staging for the imageStore epilogue, mirroring the fp16 kernel: -// SG_GRID_Y bands of MMA_M rows, each WG_TILE_N wide, row-major. A full -// WG_TILE_M x WG_TILE_N buffer would cost SG_GRID_Y/MMAS_PER_SG_M x more LDS -// and wreck occupancy. float16_t-typed because coopMatStore needs it. -const uint CSH_ROWS = SG_GRID_Y * MMA_M; -shared float16_t Csh_out[CSH_ROWS * WG_TILE_N]; -#endif - -// Running fp32 accumulator (across all groups). -coopmat - result[MMAS_PER_SG_M][MMAS_PER_SG_N]; - -// Per-group int32 MMA accumulator. -coopmat - accum_int32[MMAS_PER_SG_M][MMAS_PER_SG_N]; - - -// Byte-parallel int4 -> int8 widening. dbuf4zpg's, unchanged (B-side only). -// -// The four nibbles this shader needs from one packed uint are ALREADY one per -// byte (bits 3:0 of each byte for parity 0, bits 7:4 for parity 1), so all four -// can be widened at once instead of with a per-nibble -// shift/mask/bias-subtract/mask chain. -// -// For v in [0,15] the biased value is v-8. `v ^ 8` is exactly the 4-bit -// two's-complement pattern of v-8, because -8 == +8 (mod 16): -// v=0 -> 0x8 -> -8 v=7 -> 0xF -> -1 -// v=8 -> 0x0 -> 0 v=15 -> 0x7 -> +7 -// so the only remaining work is sign-extending bit 3 into bits 7:4 per byte. -// `sgn * 0x1E` does that with no cross-byte carry: 0x08 * 0x1E == 0xF0 exactly, -// and sgn is at most 0x08080808 so the product is at most 0xF0F0F0F0. -// -// A naive `nib - 0x08080808` would NOT work -- it borrows across byte lanes -// whenever a nibble is < 8. Shifts must be on uint, not int, so the >> is -// logical rather than arithmetic. -// -// ~5 ops per 4 weights vs ~22 for the per-nibble chain; bit-identical output. -uint widen_nibbles(const uint w, const uint parity) { - const uint nib = (parity == 0u) ? (w & 0x0F0F0F0Fu) : ((w >> 4u) & 0x0F0F0F0Fu); - const uint p = nib ^ 0x08080808u; - const uint sgn = p & 0x08080808u; - return p | (sgn * 0x1Eu); -} - -void main() { - const uvec2 tileID = uvec2(gl_WorkGroupID.xy); - const uvec2 warpInTile = uvec2( - gl_SubgroupID % SG_GRID_X, - gl_SubgroupID / SG_GRID_X); - - const uint K = uint(input_sizes.x); - const uint N = uint(output_sizes.x); - const uint N4 = (N + 3u) / 4u; - const uint nblocks_x_A = (K + 3u) >> 2u; - // A row stride in INT8 elements (dbuf4tr's binding is row-major int8, not - // the 4h4w ivec4 block layout dbuf4zpg used -- so A addressing below is in - // int8 elements, not int, and derived from nblocks_x_A so it matches the - // packer's `m_row * K4 + k4` addressing exactly (K % 4 == 0 makes them - // equal to K directly). - const uint a_row_stride_i8 = nblocks_x_A * 4u; - -#ifdef WEIGHT_INT4 - const uint num_groups = uint(num_groups_arg); - const uint CHUNKS_PER_GROUP = uint(K4_per_group) * 4u / WG_TILE_K; -#else - // Per-channel: a single quant "group" spanning all of K. The nested - // groups x chunks loop below collapses to a flat chunk loop, the wsum/wsc - // ping-pong never crosses a boundary, and the epilog runs exactly once. - const uint num_groups = 1u; - const uint CHUNKS_PER_GROUP = uint(num_groups_arg); -#endif - const uint num_chunks = num_groups * CHUNKS_PER_GROUP; - - const uint tile_m_start = WG_TILE_M * tileID.y; - const uint tile_n_start = WG_TILE_N * tileID.x; - - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - result[i][j] = coopmat(0.0); - accum_int32[i][j] = coopmat(0); - } - } - - const uint K_BLOCKS_PER_CHUNK = WG_TILE_K >> 2u; - - // --- A staging tile map (dbuf4tr's, unmodified): one MMA_M x MMA_K coopmat - // tile per subgroup per slot, dealt round-robin across the - // NUM_SUBGROUPS subgroups so every subgroup participates. Replaces - // dbuf4zpg's per-thread (m4, k4) map / a_active guard entirely -- see - // design.md D3 for why this is reused as-is, not re-derived. - const uint A_TILES_M = WG_TILE_M / MMA_M; - const uint A_TILES_K = WG_TILE_K / MMA_K; // == NUM_K_SLABS - const uint NUM_A_TILES = A_TILES_M * A_TILES_K; - const uint A_TILES_PER_SG = (NUM_A_TILES + NUM_SUBGROUPS - 1u) / NUM_SUBGROUPS; - -#ifdef WEIGHT_INT4 - // --- B staging thread map: (block, col) slots; each slot extracts one - // ColumnMajor LDS uint (4 K-contiguous sign-extended int8) --- - const uint B_TOTAL_SLOTS = K_BLOCKS_PER_CHUNK * WG_TILE_N; - const uint B_SLOTS_PER_THREAD = B_TOTAL_SLOTS / WG_SIZE; - const uint N8_PER_TILE = WG_TILE_N >> 3u; -#else - // --- B staging thread map: one (k4, n4) ivec4 block per active thread --- - const uint B_FETCH_SLOTS = K_BLOCKS_PER_CHUNK * (WG_TILE_N >> 2u); - const uint N4_PER_TILE = WG_TILE_N >> 2u; - const uint b_k4_in_chunk = gl_LocalInvocationID.x / N4_PER_TILE; - const uint b_n_uint_col = gl_LocalInvocationID.x % N4_PER_TILE; - const bool b_active = gl_LocalInvocationID.x < B_FETCH_SLOTS; -#endif - - // ===== INTERVENTION F: hoist loop-invariant B staging index math ===== - // dbuf4zpg's, unchanged -- B staging is untouched by this file's A-staging - // swap. See dbuf4zpg's header for the full rationale (ablation-attributed - // -16.8% block, mostly index arithmetic not memory traffic). -#ifdef WEIGHT_INT4 - // bv4 ownership: thread t owns one whole (slab, col) column-slab when - // t < BSH_SLICE_V4, and nothing otherwise. comp/par depend only on the - // column, so a thread's 4 k4 fetches share them and differ only in k4 -- - // 4 consecutive k4 blocks of the SAME n8 texel column. Ownership is - // subgroup-aligned (WG_TILE_N and BSH_SLICE_V4 are multiples of the subgroup - // width for every shipped geometry), so a non-owning subgroup skips the whole - // staging block on a wave-uniform branch rather than diverging. - const bool b_owns = gl_LocalInvocationID.x < BSH_SLICE_V4; - const uint b_slab = gl_LocalInvocationID.x / WG_TILE_N; - const uint b_col = gl_LocalInvocationID.x % WG_TILE_N; - const uint b_rem = b_col & 7u; - const uint b_comp1 = b_rem & 3u; - const uint b_par1 = b_rem >> 2u; - const uint b_n8blk1 = (tile_n_start >> 3u) + (b_col >> 3u); - const uint b_k4base = b_slab * (MMA_K >> 2u); - const uint b_v4_off = b_slab * WG_TILE_N + b_col; -#endif - - // Prefetch temp registers. temp_A is a coopmat array (dbuf4tr's A-staging - // technique); indices into it are [[unroll]]-resolved compile-time - // constants, never dynamic -- dynamic indexing of a coopmat array is - // exactly the construct the Xclipse/AMD-PAL compiler has miscompiled - // before. - coopmat - temp_A[A_TILES_PER_SG]; -#ifdef WEIGHT_INT4 - ivec4 temp_B[MMA_K >> 2u]; // bv4: one per k4 of the owned column-slab - float temp_wsc; -#else - ivec4 temp_B; -#endif - - // ========================================================= - // PROLOGUE - // ========================================================= - if (gl_LocalInvocationID.x < (WG_TILE_M >> 2u)) { - const uint m4 = (tile_m_start >> 2u) + gl_LocalInvocationID.x; - const vec4 sc = vec4(texelFetch(t_int8_input_scales, ivec3(m4, 0, 0), 0)); - const ivec4 zp = texelFetch(t_int8_input_zps, ivec3(m4, 0, 0), 0); - const uint base = gl_LocalInvocationID.x * 4u; - ifs_sh[base + 0u] = sc.x; ifs_sh[base + 1u] = sc.y; - ifs_sh[base + 2u] = sc.z; ifs_sh[base + 3u] = sc.w; - izp_sh[base + 0u] = zp.x; izp_sh[base + 1u] = zp.y; - izp_sh[base + 2u] = zp.z; izp_sh[base + 3u] = zp.w; - } - // Group 0 weight scales -> slice 0, and the hoisted weight-side correction - // SUM_g wsc[g][n]*wsum[g][n] accumulated across ALL groups. dbuf4zpg's - // zp-hoist, unchanged. - if (gl_LocalInvocationID.x < WG_TILE_N) { - const uint n_idx = tile_n_start + gl_LocalInvocationID.x; - f16vec4 sv0 = t_weight_scales[n_idx >> 2u]; - wsc_sh[gl_LocalInvocationID.x] = float(sv0[n_idx & 3u]); - - float corr = 0.0; - for (uint g = 0; g < num_groups; ++g) { - f16vec4 sv = t_weight_scales[g * N4 + (n_idx >> 2u)]; - corr += float(sv[n_idx & 3u]) * float(t_weight_sums[g * N + n_idx]); - } - wcorr_sh[gl_LocalInvocationID.x] = corr; - } - memoryBarrierShared(); - barrier(); - - // NOTE: dbuf4zpg builds izp_bcast/ifs_bcast AFTER the group loop, not here - // -- that is the register-pressure saving zp-hoist buys. Unchanged. - - // dbuf4: prefetch chunk 0 into temp registers, THEN store to slice 0 (no - // barrier here -- the main loop's first iteration barriers before - // reading slice 0). - // - // A staging (dbuf4tr's technique): per-subgroup coopMatLoad straight from - // the row-major global buffer. - [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { - const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; - if (t < NUM_A_TILES) { - const uint tm = t / A_TILES_K; - const uint tk = t % A_TILES_K; - coopMatLoad( - temp_A[s], t_packed_int8_input, - (tile_m_start + tm * MMA_M) * a_row_stride_i8 + tk * MMA_K, - a_row_stride_i8, - gl_CooperativeMatrixLayoutRowMajor); - } - } -#ifdef WEIGHT_INT4 - if (b_owns) { - [[unroll]] for (uint kk = 0; kk < (MMA_K >> 2u); ++kk) { -#ifdef WEIGHT_BUFFER - temp_B[kk] = t_packed_weight[(b_n8blk1 * nblocks_x_A) + b_k4base + kk]; -#else - temp_B[kk] = texelFetch(t_packed_weight, ivec2(b_k4base + kk, b_n8blk1), 0); -#endif - } - } -#else - if (b_active) { - const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; -#ifdef WEIGHT_BUFFER - temp_B = t_packed_weight[(b_k4_in_chunk * N4) + block_x_w]; -#else - temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, b_k4_in_chunk), 0); -#endif - } -#endif - { - // store chunk 0 -> slice 0 - // A staging (dbuf4tr's technique): coopMatStore into the same Ash_int8 - // slot layout dbuf4zpg's scalar scatter used to write. - [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { - const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; - if (t < NUM_A_TILES) { - const uint tm = t / A_TILES_K; - const uint tk = t % A_TILES_K; - coopMatStore( - temp_A[s], Ash_int8, - tk * A_SLAB_U32 + (tm * MMA_M) * A_STRIDE_U32, - A_STRIDE_U32, - gl_CooperativeMatrixLayoutRowMajor); - } - } -#ifdef WEIGHT_INT4 - if (b_owns) { - // ONE b128 store replaces the 4 scalar ds_write_b32 this column-slab - // previously took (spread across 4 threads). - Bsh_v4[ b_v4_off] = uvec4( - widen_nibbles(uint(temp_B[0][b_comp1]), b_par1), - widen_nibbles(uint(temp_B[1][b_comp1]), b_par1), - widen_nibbles(uint(temp_B[2][b_comp1]), b_par1), - widen_nibbles(uint(temp_B[3][b_comp1]), b_par1)); - } -#else - if (b_active) { - const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); - const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); - const uint n_col_base = b_n_uint_col * 4u; - [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { - // bv4: identical address, expressed as (uvec4 element, component). - // This INT8 path writes 4 DIFFERENT columns for one k4, so it cannot - // use a wide store -- it keeps scalar component writes. - Bsh_v4[slab_idx * WG_TILE_N + (n_col_base + n_in_blk)][k4_in_slab] = - uint(temp_B[n_in_blk]); - } - } -#endif - } - - // ========================================================= - // MAIN LOOP (dbuf4) — nested groups x chunks (kept nested; flattening it - // with a conditional coopmat epilog crashes the Xclipse PAL compiler at - // large spec-resolved trip counts). One barrier per chunk. Chunk - // iteration (global index `chunk`): - // 1. barrier — A/B slice (chunk%2) fully written; on the first chunk - // of group g, wsc slice (g%2) is too. - // 2. prefetch — chunk+1 (A tiles, B blocks) into temp; when chunk+1 - // starts a new group, also its wsc element. Skipped - // entirely on the final chunk. - // 3. int8 MMA — on slice (chunk%2) into accum_int32. - // 4. store — temp -> A/B slice ((chunk+1)%2), unpacking the weight; - // on a group boundary, wsc -> slice ((g+1)%2). - // The group epilog runs unconditionally at the tail of each group. - // ========================================================= - uint chunk = 0; - for (uint group_i = 0; group_i < num_groups; ++group_i) { - for (uint inner = 0; inner < CHUNKS_PER_GROUP; ++inner, ++chunk) { - const bool has_next = chunk + 1u < num_chunks; - const bool group_crossing = has_next && (inner + 1u == CHUNKS_PER_GROUP); - const uint cur_a = (chunk % 2u) * ASH_SLICE_U32; - const uint cur_b = (chunk % 2u) * BSH_SLICE_U32; - const uint cur_b_v4 = (chunk % 2u) * BSH_SLICE_V4; - const uint nxt_a = ((chunk + 1u) % 2u) * ASH_SLICE_U32; - const uint nxt_b = ((chunk + 1u) % 2u) * BSH_SLICE_V4; - - // coopmat-lds-fence: barrier() alone does NOT order shared stores against a - // subsequent coopMatLoad on the M51 Xclipse/AMD-PAL driver -- symptom is one - // stale MMA_M-row band of A, all columns, ~2.5% of runs, no crash (observed - // 2026-09-02 in sdpa_compute_out_coopmat.glsl). Measured cost of the fence: - // none (see this change's results). See memory - // `coopmat-lds-needs-explicit-memorybarriershared`. - memoryBarrierShared(); - barrier(); - - // --- 2. prefetch chunk+1 -> temp --- - if (has_next) { - const uint chunkK_nxt = (chunk + 1u) * WG_TILE_K; - // A staging (dbuf4tr's technique): coopMatLoad straight from global. - [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { - const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; - if (t < NUM_A_TILES) { - const uint tm = t / A_TILES_K; - const uint tk = t % A_TILES_K; - coopMatLoad( - temp_A[s], t_packed_int8_input, - (tile_m_start + tm * MMA_M) * a_row_stride_i8 + chunkK_nxt + - tk * MMA_K, - a_row_stride_i8, - gl_CooperativeMatrixLayoutRowMajor); - } - } -#ifdef WEIGHT_INT4 - if (b_owns) { - [[unroll]] for (uint kk = 0; kk < (MMA_K >> 2u); ++kk) { - const uint k4_blk = (chunkK_nxt >> 2u) + b_k4base + kk; -#ifdef WEIGHT_BUFFER - temp_B[kk] = t_packed_weight[(b_n8blk1 * nblocks_x_A) + k4_blk]; -#else - temp_B[kk] = texelFetch(t_packed_weight, ivec2(k4_blk, b_n8blk1), 0); -#endif - } - } - if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { - const uint n_idx = tile_n_start + gl_LocalInvocationID.x; - f16vec4 sv = t_weight_scales[(group_i + 1u) * N4 + (n_idx >> 2u)]; - temp_wsc = float(sv[n_idx & 3u]); - } -#else - if (b_active) { - const uint block_y_w = (chunkK_nxt >> 2u) + b_k4_in_chunk; - const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; -#ifdef WEIGHT_BUFFER - temp_B = t_packed_weight[(block_y_w * N4) + block_x_w]; -#else - temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, block_y_w), 0); -#endif - } -#endif - } - - // --- 3. int8 MMA on the cur slice --- - [[unroll]] for (uint k = 0; k < NUM_K_SLABS; ++k) { - const uint slab_a_base_u32 = cur_a + k * A_SLAB_U32; - const uint slab_b_base_u32 = cur_b + k * B_SLAB_U32; - - coopmat matA[MMAS_PER_SG_M]; - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - const uint row_a = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); - coopMatLoad( - matA[i], Ash_int8, - slab_a_base_u32 + row_a * A_STRIDE_U32, - A_STRIDE_U32, - gl_CooperativeMatrixLayoutRowMajor); - } - - coopmat matB; - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint col_b = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopMatLoad( - matB, Bsh_v4, - cur_b_v4 + k * WG_TILE_N + col_b, - 1u, - gl_CooperativeMatrixLayoutColumnMajor); - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - accum_int32[i][j] = coopMatMulAdd(matA[i], matB, accum_int32[i][j]); - } - } - } - - // --- 4. store temp (chunk+1) -> nxt slice --- - if (has_next) { - // A staging (dbuf4tr's technique): coopMatStore into the nxt slice. - [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { - const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; - if (t < NUM_A_TILES) { - const uint tm = t / A_TILES_K; - const uint tk = t % A_TILES_K; - coopMatStore( - temp_A[s], Ash_int8, - nxt_a + tk * A_SLAB_U32 + (tm * MMA_M) * A_STRIDE_U32, - A_STRIDE_U32, - gl_CooperativeMatrixLayoutRowMajor); - } - } -#ifdef WEIGHT_INT4 - if (b_owns) { - // ONE b128 store replaces the 4 scalar ds_write_b32 this column-slab - // previously took (spread across 4 threads). - Bsh_v4[nxt_b + b_v4_off] = uvec4( - widen_nibbles(uint(temp_B[0][b_comp1]), b_par1), - widen_nibbles(uint(temp_B[1][b_comp1]), b_par1), - widen_nibbles(uint(temp_B[2][b_comp1]), b_par1), - widen_nibbles(uint(temp_B[3][b_comp1]), b_par1)); - } - if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { - const uint wbase_nxt = ((group_i + 1u) % 2u) * WG_TILE_N; - wsc_sh[wbase_nxt + gl_LocalInvocationID.x] = temp_wsc; - } -#else - if (b_active) { - const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); - const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); - const uint n_col_base = b_n_uint_col * 4u; - [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { - Bsh_v4[nxt_b + slab_idx * WG_TILE_N + (n_col_base + n_in_blk)][k4_in_slab] = - uint(temp_B[n_in_blk]); - } - } -#endif - } - } // chunks - - // --- Group epilog: scale-only accumulate, reset accum --- - // dbuf4zpg's, unchanged. Just result += float(acc) * wsc. The - // zero-point subtract and the ifs multiply are hoisted out of the group - // loop (applied once below). - { - const uint wbase = (group_i % 2u) * WG_TILE_N; - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - - coopmat wsc_bcast; - coopMatLoad( - wsc_bcast, wsc_sh, - wbase + local_n_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutRowMajor); - - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - result[i][j] += - coopmat( - accum_int32[i][j]) * wsc_bcast; - accum_int32[i][j] = coopmat(0); - } - } - } - } // groups - - // --- Hoisted correction, applied ONCE: -------------------------------- - // result = ifs * ( result - izp * SUM_g wsc_g*wsum_g ) - // dbuf4zpg's, unchanged. izp/ifs are loaded here rather than before the - // group loop so they are not live across it. - { - coopmat - izpf_bcast[MMAS_PER_SG_M]; - coopmat - ifs_bcast[MMAS_PER_SG_M]; - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - const uint local_m_base = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); - coopmat izp_i; - coopMatLoad( - izp_i, izp_sh, local_m_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutColumnMajor); - izpf_bcast[i] = - coopmat(izp_i); - coopMatLoad( - ifs_bcast[i], ifs_sh, local_m_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutColumnMajor); - } - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopmat wcorr_bcast; - coopMatLoad( - wcorr_bcast, wcorr_sh, local_n_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutRowMajor); - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - result[i][j] = ifs_bcast[i] * (result[i][j] - izpf_bcast[i] * wcorr_bcast); - } - } - } - - // --- Bias (optional) --- -#ifdef HAS_BIAS - if (apply_bias > 0) { - for (uint t = gl_LocalInvocationID.x; t < WG_TILE_N; t += WG_SIZE) { - bias_sh[t] = float(t_bias[tile_n_start + t]); - } - memoryBarrierShared(); - barrier(); - } -#endif - - // --- Store result tile --- - // N for the store address math MUST come from the spec constant, not the - // sizes UBO (see out_N_arg above). -#ifdef IO_TEXTURE - // Epilogue iteration i drains accumulator row-block i from EVERY subgroup - // into Csh_out at once, so the SG_GRID_Y bands it holds are disjoint global - // row ranges; the whole workgroup then imageStores them. lr / MMA_M is the - // writing subgroup's warpInTile.y, so the global row reproduces the buffer - // path's gi exactly. - // - // PORTABILITY NOTE: the barrier() in the loop body keeps this loop rolled - // despite [[unroll]], so result[i][j] IS dynamically indexed. Coopmat arrays - // are opaque per-lane storage and dynamic indexing is exactly the construct - // the Xclipse/AMD-PAL compiler has broken before -- check this first if the - // texture variants miscompile on M51. - const uint CSH_TEXELS_PER_ROW = WG_TILE_N / 4u; - const uint CSH_TEXELS = CSH_ROWS * CSH_TEXELS_PER_ROW; - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - // Guards Csh_out against the previous iteration's readers. Inert on i == 0 - // but must stay unconditional to remain workgroup-uniform. - // coopmat-lds-fence: barrier() alone does NOT order shared stores against a - // subsequent coopMatLoad on the M51 Xclipse/AMD-PAL driver -- symptom is one - // stale MMA_M-row band of A, all columns, ~2.5% of runs, no crash (observed - // 2026-09-02 in sdpa_compute_out_coopmat.glsl). Measured cost of the fence: - // none (see this change's results). See memory - // `coopmat-lds-needs-explicit-memorybarriershared`. - memoryBarrierShared(); - barrier(); - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { -#ifdef HAS_BIAS - if (apply_bias > 0) { - const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopmat bias_tile; - coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); - result[i][j] += bias_tile; - } -#endif - coopmat out_tile = - coopmat(result[i][j]); - coopMatStore( - out_tile, Csh_out, - warpInTile.y * MMA_M * WG_TILE_N + - MMA_N * (MMAS_PER_SG_N * warpInTile.x + j), - WG_TILE_N, - gl_CooperativeMatrixLayoutRowMajor); - } - memoryBarrierShared(); - barrier(); - - for (uint t = gl_LocalInvocationID.x; t < CSH_TEXELS; t += WG_SIZE) { - const uint lr = t / CSH_TEXELS_PER_ROW; - const uint lc4 = t % CSH_TEXELS_PER_ROW; - const uint m = - tile_m_start + (lr / MMA_M) * SG_TILE_M + i * MMA_M + (lr % MMA_M); - const uint base = lr * WG_TILE_N + lc4 * 4u; - imageStore( - t_output, - ivec3(tile_n_start / 4u + lc4, m, 0), - vec4( - float(Csh_out[base]), - float(Csh_out[base + 1u]), - float(Csh_out[base + 2u]), - float(Csh_out[base + 3u]))); - } - } -#else - const uint N_out = uint(out_N_arg); - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint gi = tile_m_start + MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); - const uint gj = tile_n_start + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - -#ifdef HAS_BIAS - if (apply_bias > 0) { - const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopmat bias_tile; - coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); - result[i][j] += bias_tile; - } -#endif - - coopmat out_tile = - coopmat(result[i][j]); - coopMatStore( - out_tile, t_output, - gi * N_out + gj, N_out, - gl_CooperativeMatrixLayoutRowMajor); - } - } -#endif // IO_TEXTURE -} diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbv4.yaml b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbv4.yaml deleted file mode 100644 index c1ca672d347..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbv4.yaml +++ /dev/null @@ -1,68 +0,0 @@ -# !!! KNOWN-INCORRECT -- DO NOT PROMOTE, DO NOT TIME !!! -# Fails the correctness gate deterministically (3/3 reps) on exactly the -# num_groups == 2 shapes (K=256: M128_K256_N128, M256_K256_N256). Every K=128, -# K=2048 and K=4096 case PASSES. Root cause NOT yet identified: the -# (slab, col) -> (n8blk, k4, component, parity) mapping was re-derived and is -# provably identical to the baseline slot map, so the defect is in something -# adjacent to the ownership change, not in the address arithmetic. -# Kept opt-in only (never a default) as a recorded negative result. -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -# "zpg" + "-tr" combination: tsweep_dbuf4zpg with its per-thread scalar -# A-staging replaced by tsweep_dbuf4tr's coopMat-mediated A staging. B -# staging, zp-hoist, byte-parallel nibble widening are all dbuf4zpg's, -# unchanged -- only A-staging differs. Requires t_packed_int8_input in the -# ROW-MAJOR kPackedInt8_4W layout (same requirement as tsweep_dbuf4tr). -# -# PROMOTED 2026-09-01 as the shipped default -- see QuantizedLinear.cpp's -# dq8ca_coopmat_variant() for the full validation record (10/10 buffer, 6/6 -# texture3d correctness across 1B/3B/8B; +4.2% real e2e prefill; 46.50% -> -# 49.94% efficiency on 8B). Also selectable explicitly via -# ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4zpgtr_txkgs. -# See openspec/changes/archive/2026-08-31-dq8ca-tr-staged-a-on-zpg. -# -# Single seed tile: the prior dbuf4zpg default's own tile (t128x64k32g42s32). -# A re-sweep against this shader's own (lower) register-pressure profile was -# run as a follow-up (coopmat-tr-tilesweep-4w-port) and found no better tile -# -- this remains the best known geometry. - -linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbv4: - parameter_names_with_default_values: - PRECISION: highp - HAS_BIAS: false - IO_STORAGE: buffer - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - MMA_M: 16 - MMA_N: 16 - MMA_K: 16 - WG_TILE_M: 64 - WG_TILE_N: 32 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - shader_variants: - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbv4_t128x64k32g42s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbv4_t128x64k32g42s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw.glsl b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw.glsl deleted file mode 100644 index 2331b04734f..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw.glsl +++ /dev/null @@ -1,781 +0,0 @@ -/* - * Copyright (c) Meta Platforms, Inc. and affiliates. - * All rights reserved. - * - * This source code is licensed under the BSD-style license found in the - * LICENSE file in the root directory of this source tree. - */ - -/* - * "zpg" + "-tr" combination: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpg.glsl - * with its per-thread scalar A-staging replaced by - * linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr.glsl's coopMat-mediated A - * staging, PLUS "bw": B staging ownership moved from (block,col) slots spread - * over 4 threads to ONE thread per (K-slab, column), writing its 4 uints at - * consecutive LDS indices so they are merge-eligible. Shared-array types, the - * compute-side loads, the quantization math and every safeguard are - * dbuf4zpgtr's, unchanged. This is - * an ADDITIVE combination, not a redesign: every non-A-staging block below - * (B staging: coalesced write, no skew; zp-hoist: izp/ifs applied once after - * the group loop via wcorr_sh; byte-parallel nibble widening; static - * A_ALWAYS_ACTIVE branch elision -- N/A here, see below; group epilog; - * bias/store epilogue) is byte-identical to dbuf4zpg's. Only the A-staging - * block (prologue load+store, main-loop prefetch+store) is dbuf4tr's, - * verbatim. - * - * dbuf4zpg's per-thread A staging used an `a_active` guard (statically always - * true when A_ACTIVE_THREADS == WG_SIZE, via the A_MAP_FULL-gated - * A_ALWAYS_ACTIVE macro). dbuf4tr's per-SUBGROUP tile map has no equivalent - * concept -- every subgroup participates via a `t < NUM_A_TILES` guard that - * depends only on gl_SubgroupID, not gl_LocalInvocationID.x -- so - * A_MAP_FULL/A_ALWAYS_ACTIVE is dropped entirely in this file; it would be - * dead code for the new A-staging block. - * - * Rationale for combining this way (not the reverse) and why this is worth - * building at all: see this change's design.md D0-D3. In short -- the only - * existing measurement of dbuf4tr's A-staging technique (28.72-30.51%, - * dq8ca-arch-redesign) was taken against dbuf4tr's own pre-zpg baseline - * (old B skew, no byte-parallel widening, no branch elision) -- a materially - * weaker shader than the 46.49-46.50% dbuf4zpg this file now combines it - * with. This file exists to answer whether that combination performs - * differently now that register pressure is already reduced. - * - * A staging (the actual delta from dbuf4zpg): - * dbuf4zpg: per-thread (m4, k4) ivec4 fetch, hoisted a_lds_off0/a_glb_row; - * only A_ACTIVE_THREADS invocations participate, each scattering - * 4 rows into Ash_int8 with 4 scalar stores. - * this file: per-SUBGROUP MMA_M x MMA_K tile fetch via coopMatLoad straight - * from a ROW-MAJOR (kPackedInt8_4W) int8 activation buffer, then - * coopMatStore into the same Ash_int8 slot -- dbuf4tr's mapping, - * unmodified (not re-derived; see design.md D3). - * - * t_packed_int8_input is therefore bound the same way dbuf4tr binds it: a - * SCALAR int8_t array in the kPackedInt8_4W layout (plain row-major int8, - * row stride K), produced by quantize_and_pack_4w_with_group_sums.glsl. - * QuantizedLinear.cpp's dq8ca_variant_wants_rowmajor_a() must recognize this - * file's variant token (tsweep_dbuf4zpgtr_t...) the same way it already - * recognizes tsweep_dbuf4tr_t/trm_t/trd_t, so graph-build time (packer - * selection) and dispatch time (kernel selection) cannot disagree. - * - * B CANNOT be coopmat-staged (int4 nibble unpack; a coopmat's per-lane layout - * is opaque to hand-assembly from unpacked registers) -- unchanged from both - * parent files. B staging below is dbuf4zpg's byte-parallel, coalesced, - * no-skew version, untouched. - * - * The loop structure is dbuf4's (both parents share it), unchanged: - * prologue: prefetch chunk 0 -> temp, store to slice 0 (no barrier) - * per iter: barrier -> prefetch(next) -> MMA(cur) -> store(next) - * kept nested (groups x chunks) with an unconditional group epilog -- - * flattening it crashes the Xclipse PAL compiler at large spec-resolved trip - * counts (see dbuf2's own header). - * - * Selected via - * ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4zpgtr_txkgs<32|64> - * (QuantizedLinear.cpp), additive to the tsweep_dbuf4zpg_t..., tsweep_dbuf4tr_t... - * and tsweep_t... namespaces. NOT the default -- unvalidated until it passes - * repeated test_llama_microbench --correctness-only runs (see - * dq8ca_coopmat_variant()'s comment on why a single pass is not proof). - * - * Performs: out[M,N] = dequant(int8_act) * dequant(int_w) (+ bias) - * via coopmat x coopmat -> coopmat on the matrix unit. - * - * Hard preconditions (dbuf4zpg's, plus dbuf4tr's row-major/alignment ones): - * M % WG_TILE_M == 0, N % WG_TILE_N == 0, K % WG_TILE_K == 0, - * group_size % WG_TILE_K == 0, K % 4 == 0, - * WG_TILE_M % MMA_M == 0, WG_TILE_K % MMA_K == 0, - * t_packed_int8_input in kPackedInt8_4W (row-major) layout, - * device exposes coopmatx-> at 16x16x16. - */ - -#version 450 core - -#extension GL_KHR_cooperative_matrix : require -#extension GL_KHR_memory_scope_semantics : require -#extension GL_KHR_shader_subgroup_basic : enable -#extension GL_EXT_shader_explicit_arithmetic_types : require -#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require -// 8-bit SSBO access: A is bound as a scalar int8_t array so that the -// coopMatLoad below has a MATCHING component type (see dbuf4tr's header for -// why the type must match on this driver). -#extension GL_EXT_shader_8bit_storage : require -#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require -#extension GL_EXT_control_flow_attributes : enable - -#define PRECISION ${PRECISION} - -$if WEIGHT_NBITS == 4: - #define WEIGHT_INT4 - -$if HAS_BIAS: - #define HAS_BIAS - -$if WEIGHT_STORAGE == "buffer": - #define WEIGHT_BUFFER - -$if IO_STORAGE == "texture3d": - #define IO_TEXTURE - -layout(std430) buffer; - -#include "common.glslh" - -// Bindings — match add_linear_dqa_qw_node arg order: -// output(0), fp_input(1), packed_int8_input(2), int_input_sums(3 - unused), -// input_scales(4), input_zps(5), packed_weight(6), weight_sums(7), -// weight_scales(8), bias(9). -${layout_declare_tensor(B, "w", "t_output", "half", IO_STORAGE, is_scalar_array=True)} -// t_input is unread here -- the activations arrive already quantized in -// t_packed_int8_input -- but stays declared so the binding layout matches the -// dispatch site. It tracks IO_STORAGE so the two IO tensors stay consistent. -${layout_declare_tensor(B, "r", "t_input", "half", IO_STORAGE, is_scalar_array=False)} -// ROW-MAJOR (kPackedInt8_4W) packed activations, bound as a scalar int8_t -// array (row stride = K int8) -- dbuf4tr's binding, unchanged. The stock -// 4h4w layout dbuf4zpg uses is NOT row-major (component index selects a row, -// non-affine), so it cannot be addressed by any coopMatLoad. -${layout_declare_tensor(B, "r", "t_packed_int8_input", "int8", "buffer", is_scalar_array=True)} -${layout_declare_tensor(B, "r", "t_int8_input_sums", "int", "buffer", is_scalar_array=True)} -${layout_declare_tensor(B, "r", "t_int8_input_scales", "half", "texture3d")} -${layout_declare_tensor(B, "r", "t_int8_input_zps", "int8", "texture3d")} -${layout_declare_tensor(B, "r", "t_packed_weight", "int", WEIGHT_STORAGE, is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_weight_sums", "int", "buffer", is_scalar_array=True)} -${layout_declare_tensor(B, "r", "t_weight_scales", "half", "buffer", is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_bias", "half", "buffer", is_scalar_array=True)} - -${layout_declare_ubo(B, "ivec4", "output_sizes")} -${layout_declare_ubo(B, "ivec4", "input_sizes")} - -layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; - -${layout_declare_spec_const(C, "int", "apply_bias", "0")} -// INT4 only; inert (0) for INT8 so the dispatcher's spec list lines up. -${layout_declare_spec_const(C, "int", "K4_per_group", "0")} -${layout_declare_spec_const(C, "int", "num_groups_arg", "0")} -${layout_declare_spec_const(C, "int", "out_N_arg", "0")} - -// Tile geometry -const uint MMA_M = ${MMA_M}; -const uint MMA_N = ${MMA_N}; -const uint MMA_K = ${MMA_K}; - -const uint WG_TILE_M = ${WG_TILE_M}; -const uint WG_TILE_N = ${WG_TILE_N}; -const uint WG_TILE_K = ${WG_TILE_K}; - -const uint SG_GRID_X = ${SG_GRID_X}; -const uint SG_GRID_Y = ${SG_GRID_Y}; -const uint SUBGROUP_SIZE = ${SUBGROUP_SIZE}; -const uint NUM_SUBGROUPS = SG_GRID_X * SG_GRID_Y; -const uint WG_SIZE = NUM_SUBGROUPS * SUBGROUP_SIZE; - -const uint SG_TILE_M = WG_TILE_M / SG_GRID_Y; -const uint SG_TILE_N = WG_TILE_N / SG_GRID_X; -const uint MMAS_PER_SG_M = SG_TILE_M / MMA_M; -const uint MMAS_PER_SG_N = SG_TILE_N / MMA_N; - -const uint A_SLAB_INT8 = WG_TILE_M * MMA_K; -const uint B_USEFUL_U32 = MMA_K / 4u; -// No skew + coalesced write -- dbuf4zpg's B fix, unchanged (this file does -// not touch B staging at all). -const uint B_STRIDE_U32 = B_USEFUL_U32; -const uint B_SLAB_U32 = WG_TILE_N * B_STRIDE_U32; -const uint NUM_K_SLABS = WG_TILE_K / MMA_K; - -const uint A_SLAB_U32 = A_SLAB_INT8 / 4u; -const uint A_STRIDE_U32 = MMA_K / 4u; - -// One ping-pong slice covers all K-slabs of one chunk. -const uint ASH_SLICE_U32 = NUM_K_SLABS * A_SLAB_U32; -const uint BSH_SLICE_U32 = NUM_K_SLABS * B_SLAB_U32; - -// Double-buffered MMA operand staging. -shared uint Ash_int8[2u * ASH_SLICE_U32]; -shared uint Bsh_int8[2u * BSH_SLICE_U32]; - -// Per-WG-tile-row activation params (loaded ONCE at WG start; constant -// across groups). -shared int izp_sh[WG_TILE_M]; // int32 (cast from int8 source) for broadcast -shared float ifs_sh[WG_TILE_M]; // float32 (cast from fp16 source) for broadcast - -// Per-(group, output-channel) weight params, ping-ponged by group parity. -// (For per-channel INT8 only slice 0 is ever used.) -shared float wsc_sh[2u * WG_TILE_N]; -// SUM_g wsc[g][n]*wsum[g][n] per output channel -- weight-side only, so it is -// accumulated once in the prologue. dbuf4zpg's zp-hoist, unchanged. -shared float wcorr_sh[WG_TILE_N]; - -#ifdef HAS_BIAS -shared float bias_sh[WG_TILE_N]; -#endif - -#ifdef IO_TEXTURE -// Result staging for the imageStore epilogue, mirroring the fp16 kernel: -// SG_GRID_Y bands of MMA_M rows, each WG_TILE_N wide, row-major. A full -// WG_TILE_M x WG_TILE_N buffer would cost SG_GRID_Y/MMAS_PER_SG_M x more LDS -// and wreck occupancy. float16_t-typed because coopMatStore needs it. -const uint CSH_ROWS = SG_GRID_Y * MMA_M; -shared float16_t Csh_out[CSH_ROWS * WG_TILE_N]; -#endif - -// Running fp32 accumulator (across all groups). -coopmat - result[MMAS_PER_SG_M][MMAS_PER_SG_N]; - -// Per-group int32 MMA accumulator. -coopmat - accum_int32[MMAS_PER_SG_M][MMAS_PER_SG_N]; - - -// Byte-parallel int4 -> int8 widening. dbuf4zpg's, unchanged (B-side only). -// -// The four nibbles this shader needs from one packed uint are ALREADY one per -// byte (bits 3:0 of each byte for parity 0, bits 7:4 for parity 1), so all four -// can be widened at once instead of with a per-nibble -// shift/mask/bias-subtract/mask chain. -// -// For v in [0,15] the biased value is v-8. `v ^ 8` is exactly the 4-bit -// two's-complement pattern of v-8, because -8 == +8 (mod 16): -// v=0 -> 0x8 -> -8 v=7 -> 0xF -> -1 -// v=8 -> 0x0 -> 0 v=15 -> 0x7 -> +7 -// so the only remaining work is sign-extending bit 3 into bits 7:4 per byte. -// `sgn * 0x1E` does that with no cross-byte carry: 0x08 * 0x1E == 0xF0 exactly, -// and sgn is at most 0x08080808 so the product is at most 0xF0F0F0F0. -// -// A naive `nib - 0x08080808` would NOT work -- it borrows across byte lanes -// whenever a nibble is < 8. Shifts must be on uint, not int, so the >> is -// logical rather than arithmetic. -// -// ~5 ops per 4 weights vs ~22 for the per-nibble chain; bit-identical output. -uint widen_nibbles(const uint w, const uint parity) { - const uint nib = (parity == 0u) ? (w & 0x0F0F0F0Fu) : ((w >> 4u) & 0x0F0F0F0Fu); - const uint p = nib ^ 0x08080808u; - const uint sgn = p & 0x08080808u; - return p | (sgn * 0x1Eu); -} - -void main() { - const uvec2 tileID = uvec2(gl_WorkGroupID.xy); - const uvec2 warpInTile = uvec2( - gl_SubgroupID % SG_GRID_X, - gl_SubgroupID / SG_GRID_X); - - const uint K = uint(input_sizes.x); - const uint N = uint(output_sizes.x); - const uint N4 = (N + 3u) / 4u; - const uint nblocks_x_A = (K + 3u) >> 2u; - // A row stride in INT8 elements (dbuf4tr's binding is row-major int8, not - // the 4h4w ivec4 block layout dbuf4zpg used -- so A addressing below is in - // int8 elements, not int, and derived from nblocks_x_A so it matches the - // packer's `m_row * K4 + k4` addressing exactly (K % 4 == 0 makes them - // equal to K directly). - const uint a_row_stride_i8 = nblocks_x_A * 4u; - -#ifdef WEIGHT_INT4 - const uint num_groups = uint(num_groups_arg); - const uint CHUNKS_PER_GROUP = uint(K4_per_group) * 4u / WG_TILE_K; -#else - // Per-channel: a single quant "group" spanning all of K. The nested - // groups x chunks loop below collapses to a flat chunk loop, the wsum/wsc - // ping-pong never crosses a boundary, and the epilog runs exactly once. - const uint num_groups = 1u; - const uint CHUNKS_PER_GROUP = uint(num_groups_arg); -#endif - const uint num_chunks = num_groups * CHUNKS_PER_GROUP; - - const uint tile_m_start = WG_TILE_M * tileID.y; - const uint tile_n_start = WG_TILE_N * tileID.x; - - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - result[i][j] = coopmat(0.0); - accum_int32[i][j] = coopmat(0); - } - } - - const uint K_BLOCKS_PER_CHUNK = WG_TILE_K >> 2u; - - // --- A staging tile map (dbuf4tr's, unmodified): one MMA_M x MMA_K coopmat - // tile per subgroup per slot, dealt round-robin across the - // NUM_SUBGROUPS subgroups so every subgroup participates. Replaces - // dbuf4zpg's per-thread (m4, k4) map / a_active guard entirely -- see - // design.md D3 for why this is reused as-is, not re-derived. - const uint A_TILES_M = WG_TILE_M / MMA_M; - const uint A_TILES_K = WG_TILE_K / MMA_K; // == NUM_K_SLABS - const uint NUM_A_TILES = A_TILES_M * A_TILES_K; - const uint A_TILES_PER_SG = (NUM_A_TILES + NUM_SUBGROUPS - 1u) / NUM_SUBGROUPS; - -#ifdef WEIGHT_INT4 - // --- B staging thread map: (block, col) slots; each slot extracts one - // ColumnMajor LDS uint (4 K-contiguous sign-extended int8) --- - const uint B_TOTAL_SLOTS = K_BLOCKS_PER_CHUNK * WG_TILE_N; - const uint B_SLOTS_PER_THREAD = B_TOTAL_SLOTS / WG_SIZE; - const uint N8_PER_TILE = WG_TILE_N >> 3u; -#else - // --- B staging thread map: one (k4, n4) ivec4 block per active thread --- - const uint B_FETCH_SLOTS = K_BLOCKS_PER_CHUNK * (WG_TILE_N >> 2u); - const uint N4_PER_TILE = WG_TILE_N >> 2u; - const uint b_k4_in_chunk = gl_LocalInvocationID.x / N4_PER_TILE; - const uint b_n_uint_col = gl_LocalInvocationID.x % N4_PER_TILE; - const bool b_active = gl_LocalInvocationID.x < B_FETCH_SLOTS; -#endif - - // ===== INTERVENTION F: hoist loop-invariant B staging index math ===== - // dbuf4zpg's, unchanged -- B staging is untouched by this file's A-staging - // swap. See dbuf4zpg's header for the full rationale (ablation-attributed - // -16.8% block, mostly index arithmetic not memory traffic). -#ifdef WEIGHT_INT4 - // ===== bw: SOURCE-side change only -- one thread owns a whole (slab, col) - // column-slab and writes its 4 uints at CONSECUTIVE indices, so the four - // ds_write_b32 are merge-eligible. The shared array stays `uint` (the proven - // 4x packing ratio) and the compute-side coopMatLoad is byte-for-byte - // unchanged -- deliberately NOT the uvec4 retype, which fails correctness on - // this driver (16x packing; see the 8x coopMatStore precedent). - // comp/par depend only on the column, so a thread's 4 fetches share them and - // differ only in k4: 4 consecutive k4 blocks of the same n8 texel column. - const uint B_OWNERS = NUM_K_SLABS * WG_TILE_N; - const bool b_owns = gl_LocalInvocationID.x < B_OWNERS; - const uint b_slab = gl_LocalInvocationID.x / WG_TILE_N; - const uint b_col = gl_LocalInvocationID.x % WG_TILE_N; - const uint b_rem = b_col & 7u; - const uint b_comp1 = b_rem & 3u; - const uint b_par1 = b_rem >> 2u; - const uint b_n8blk1 = (tile_n_start >> 3u) + (b_col >> 3u); - const uint b_k4base = b_slab * (MMA_K >> 2u); - // same address the destination-oriented map produced, just all 4 in one thread - const uint b_lds_base = b_slab * B_SLAB_U32 + b_col * B_STRIDE_U32; -#endif - - // Prefetch temp registers. temp_A is a coopmat array (dbuf4tr's A-staging - // technique); indices into it are [[unroll]]-resolved compile-time - // constants, never dynamic -- dynamic indexing of a coopmat array is - // exactly the construct the Xclipse/AMD-PAL compiler has miscompiled - // before. - coopmat - temp_A[A_TILES_PER_SG]; -#ifdef WEIGHT_INT4 - ivec4 temp_B[MMA_K >> 2u]; // bw: one per k4 of the owned column-slab - float temp_wsc; -#else - ivec4 temp_B; -#endif - - // ========================================================= - // PROLOGUE - // ========================================================= - if (gl_LocalInvocationID.x < (WG_TILE_M >> 2u)) { - const uint m4 = (tile_m_start >> 2u) + gl_LocalInvocationID.x; - const vec4 sc = vec4(texelFetch(t_int8_input_scales, ivec3(m4, 0, 0), 0)); - const ivec4 zp = texelFetch(t_int8_input_zps, ivec3(m4, 0, 0), 0); - const uint base = gl_LocalInvocationID.x * 4u; - ifs_sh[base + 0u] = sc.x; ifs_sh[base + 1u] = sc.y; - ifs_sh[base + 2u] = sc.z; ifs_sh[base + 3u] = sc.w; - izp_sh[base + 0u] = zp.x; izp_sh[base + 1u] = zp.y; - izp_sh[base + 2u] = zp.z; izp_sh[base + 3u] = zp.w; - } - // Group 0 weight scales -> slice 0, and the hoisted weight-side correction - // SUM_g wsc[g][n]*wsum[g][n] accumulated across ALL groups. dbuf4zpg's - // zp-hoist, unchanged. - if (gl_LocalInvocationID.x < WG_TILE_N) { - const uint n_idx = tile_n_start + gl_LocalInvocationID.x; - f16vec4 sv0 = t_weight_scales[n_idx >> 2u]; - wsc_sh[gl_LocalInvocationID.x] = float(sv0[n_idx & 3u]); - - float corr = 0.0; - for (uint g = 0; g < num_groups; ++g) { - f16vec4 sv = t_weight_scales[g * N4 + (n_idx >> 2u)]; - corr += float(sv[n_idx & 3u]) * float(t_weight_sums[g * N + n_idx]); - } - wcorr_sh[gl_LocalInvocationID.x] = corr; - } - memoryBarrierShared(); - barrier(); - - // NOTE: dbuf4zpg builds izp_bcast/ifs_bcast AFTER the group loop, not here - // -- that is the register-pressure saving zp-hoist buys. Unchanged. - - // dbuf4: prefetch chunk 0 into temp registers, THEN store to slice 0 (no - // barrier here -- the main loop's first iteration barriers before - // reading slice 0). - // - // A staging (dbuf4tr's technique): per-subgroup coopMatLoad straight from - // the row-major global buffer. - [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { - const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; - if (t < NUM_A_TILES) { - const uint tm = t / A_TILES_K; - const uint tk = t % A_TILES_K; - coopMatLoad( - temp_A[s], t_packed_int8_input, - (tile_m_start + tm * MMA_M) * a_row_stride_i8 + tk * MMA_K, - a_row_stride_i8, - gl_CooperativeMatrixLayoutRowMajor); - } - } -#ifdef WEIGHT_INT4 - if (b_owns) { - [[unroll]] for (uint kk = 0; kk < (MMA_K >> 2u); ++kk) { -#ifdef WEIGHT_BUFFER - temp_B[kk] = t_packed_weight[(b_n8blk1 * nblocks_x_A) + b_k4base + kk]; -#else - temp_B[kk] = texelFetch(t_packed_weight, ivec2(b_k4base + kk, b_n8blk1), 0); -#endif - } - } -#else - if (b_active) { - const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; -#ifdef WEIGHT_BUFFER - temp_B = t_packed_weight[(b_k4_in_chunk * N4) + block_x_w]; -#else - temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, b_k4_in_chunk), 0); -#endif - } -#endif - { - // store chunk 0 -> slice 0 - // A staging (dbuf4tr's technique): coopMatStore into the same Ash_int8 - // slot layout dbuf4zpg's scalar scatter used to write. - [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { - const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; - if (t < NUM_A_TILES) { - const uint tm = t / A_TILES_K; - const uint tk = t % A_TILES_K; - coopMatStore( - temp_A[s], Ash_int8, - tk * A_SLAB_U32 + (tm * MMA_M) * A_STRIDE_U32, - A_STRIDE_U32, - gl_CooperativeMatrixLayoutRowMajor); - } - } -#ifdef WEIGHT_INT4 - if (b_owns) { - // 4 CONSECUTIVE uints from one thread -> merge-eligible into a wider - // store, versus 4 scalar stores from 4 different threads before. - [[unroll]] for (uint kk = 0; kk < (MMA_K >> 2u); ++kk) { - Bsh_int8[ b_lds_base + kk] = - widen_nibbles(uint(temp_B[kk][b_comp1]), b_par1); - } - } -#else - if (b_active) { - const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); - const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); - const uint n_col_base = b_n_uint_col * 4u; - [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { - Bsh_int8[slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab] = - uint(temp_B[n_in_blk]); - } - } -#endif - } - - // ========================================================= - // MAIN LOOP (dbuf4) — nested groups x chunks (kept nested; flattening it - // with a conditional coopmat epilog crashes the Xclipse PAL compiler at - // large spec-resolved trip counts). One barrier per chunk. Chunk - // iteration (global index `chunk`): - // 1. barrier — A/B slice (chunk%2) fully written; on the first chunk - // of group g, wsc slice (g%2) is too. - // 2. prefetch — chunk+1 (A tiles, B blocks) into temp; when chunk+1 - // starts a new group, also its wsc element. Skipped - // entirely on the final chunk. - // 3. int8 MMA — on slice (chunk%2) into accum_int32. - // 4. store — temp -> A/B slice ((chunk+1)%2), unpacking the weight; - // on a group boundary, wsc -> slice ((g+1)%2). - // The group epilog runs unconditionally at the tail of each group. - // ========================================================= - uint chunk = 0; - for (uint group_i = 0; group_i < num_groups; ++group_i) { - for (uint inner = 0; inner < CHUNKS_PER_GROUP; ++inner, ++chunk) { - const bool has_next = chunk + 1u < num_chunks; - const bool group_crossing = has_next && (inner + 1u == CHUNKS_PER_GROUP); - const uint cur_a = (chunk % 2u) * ASH_SLICE_U32; - const uint cur_b = (chunk % 2u) * BSH_SLICE_U32; - const uint nxt_a = ((chunk + 1u) % 2u) * ASH_SLICE_U32; - const uint nxt_b = ((chunk + 1u) % 2u) * BSH_SLICE_U32; - - // coopmat-lds-fence: barrier() alone does NOT order shared stores against a - // subsequent coopMatLoad on the M51 Xclipse/AMD-PAL driver -- symptom is one - // stale MMA_M-row band of A, all columns, ~2.5% of runs, no crash (observed - // 2026-09-02 in sdpa_compute_out_coopmat.glsl). Measured cost of the fence: - // none (see this change's results). See memory - // `coopmat-lds-needs-explicit-memorybarriershared`. - memoryBarrierShared(); - barrier(); - - // --- 2. prefetch chunk+1 -> temp --- - if (has_next) { - const uint chunkK_nxt = (chunk + 1u) * WG_TILE_K; - // A staging (dbuf4tr's technique): coopMatLoad straight from global. - [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { - const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; - if (t < NUM_A_TILES) { - const uint tm = t / A_TILES_K; - const uint tk = t % A_TILES_K; - coopMatLoad( - temp_A[s], t_packed_int8_input, - (tile_m_start + tm * MMA_M) * a_row_stride_i8 + chunkK_nxt + - tk * MMA_K, - a_row_stride_i8, - gl_CooperativeMatrixLayoutRowMajor); - } - } -#ifdef WEIGHT_INT4 - if (b_owns) { - [[unroll]] for (uint kk = 0; kk < (MMA_K >> 2u); ++kk) { - const uint k4_blk = (chunkK_nxt >> 2u) + b_k4base + kk; -#ifdef WEIGHT_BUFFER - temp_B[kk] = t_packed_weight[(b_n8blk1 * nblocks_x_A) + k4_blk]; -#else - temp_B[kk] = texelFetch(t_packed_weight, ivec2(k4_blk, b_n8blk1), 0); -#endif - } - } - if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { - const uint n_idx = tile_n_start + gl_LocalInvocationID.x; - f16vec4 sv = t_weight_scales[(group_i + 1u) * N4 + (n_idx >> 2u)]; - temp_wsc = float(sv[n_idx & 3u]); - } -#else - if (b_active) { - const uint block_y_w = (chunkK_nxt >> 2u) + b_k4_in_chunk; - const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; -#ifdef WEIGHT_BUFFER - temp_B = t_packed_weight[(block_y_w * N4) + block_x_w]; -#else - temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, block_y_w), 0); -#endif - } -#endif - } - - // --- 3. int8 MMA on the cur slice --- - [[unroll]] for (uint k = 0; k < NUM_K_SLABS; ++k) { - const uint slab_a_base_u32 = cur_a + k * A_SLAB_U32; - const uint slab_b_base_u32 = cur_b + k * B_SLAB_U32; - - coopmat matA[MMAS_PER_SG_M]; - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - const uint row_a = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); - coopMatLoad( - matA[i], Ash_int8, - slab_a_base_u32 + row_a * A_STRIDE_U32, - A_STRIDE_U32, - gl_CooperativeMatrixLayoutRowMajor); - } - - coopmat matB; - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint col_b = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopMatLoad( - matB, Bsh_int8, - slab_b_base_u32 + col_b * B_STRIDE_U32, - B_STRIDE_U32, - gl_CooperativeMatrixLayoutColumnMajor); - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - accum_int32[i][j] = coopMatMulAdd(matA[i], matB, accum_int32[i][j]); - } - } - } - - // --- 4. store temp (chunk+1) -> nxt slice --- - if (has_next) { - // A staging (dbuf4tr's technique): coopMatStore into the nxt slice. - [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { - const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; - if (t < NUM_A_TILES) { - const uint tm = t / A_TILES_K; - const uint tk = t % A_TILES_K; - coopMatStore( - temp_A[s], Ash_int8, - nxt_a + tk * A_SLAB_U32 + (tm * MMA_M) * A_STRIDE_U32, - A_STRIDE_U32, - gl_CooperativeMatrixLayoutRowMajor); - } - } -#ifdef WEIGHT_INT4 - if (b_owns) { - // 4 CONSECUTIVE uints from one thread -> merge-eligible into a wider - // store, versus 4 scalar stores from 4 different threads before. - [[unroll]] for (uint kk = 0; kk < (MMA_K >> 2u); ++kk) { - Bsh_int8[nxt_b + b_lds_base + kk] = - widen_nibbles(uint(temp_B[kk][b_comp1]), b_par1); - } - } - if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { - const uint wbase_nxt = ((group_i + 1u) % 2u) * WG_TILE_N; - wsc_sh[wbase_nxt + gl_LocalInvocationID.x] = temp_wsc; - } -#else - if (b_active) { - const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); - const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); - const uint n_col_base = b_n_uint_col * 4u; - [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { - Bsh_int8[nxt_b + slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab] = - uint(temp_B[n_in_blk]); - } - } -#endif - } - } // chunks - - // --- Group epilog: scale-only accumulate, reset accum --- - // dbuf4zpg's, unchanged. Just result += float(acc) * wsc. The - // zero-point subtract and the ifs multiply are hoisted out of the group - // loop (applied once below). - { - const uint wbase = (group_i % 2u) * WG_TILE_N; - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - - coopmat wsc_bcast; - coopMatLoad( - wsc_bcast, wsc_sh, - wbase + local_n_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutRowMajor); - - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - result[i][j] += - coopmat( - accum_int32[i][j]) * wsc_bcast; - accum_int32[i][j] = coopmat(0); - } - } - } - } // groups - - // --- Hoisted correction, applied ONCE: -------------------------------- - // result = ifs * ( result - izp * SUM_g wsc_g*wsum_g ) - // dbuf4zpg's, unchanged. izp/ifs are loaded here rather than before the - // group loop so they are not live across it. - { - coopmat - izpf_bcast[MMAS_PER_SG_M]; - coopmat - ifs_bcast[MMAS_PER_SG_M]; - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - const uint local_m_base = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); - coopmat izp_i; - coopMatLoad( - izp_i, izp_sh, local_m_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutColumnMajor); - izpf_bcast[i] = - coopmat(izp_i); - coopMatLoad( - ifs_bcast[i], ifs_sh, local_m_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutColumnMajor); - } - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopmat wcorr_bcast; - coopMatLoad( - wcorr_bcast, wcorr_sh, local_n_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutRowMajor); - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - result[i][j] = ifs_bcast[i] * (result[i][j] - izpf_bcast[i] * wcorr_bcast); - } - } - } - - // --- Bias (optional) --- -#ifdef HAS_BIAS - if (apply_bias > 0) { - for (uint t = gl_LocalInvocationID.x; t < WG_TILE_N; t += WG_SIZE) { - bias_sh[t] = float(t_bias[tile_n_start + t]); - } - memoryBarrierShared(); - barrier(); - } -#endif - - // --- Store result tile --- - // N for the store address math MUST come from the spec constant, not the - // sizes UBO (see out_N_arg above). -#ifdef IO_TEXTURE - // Epilogue iteration i drains accumulator row-block i from EVERY subgroup - // into Csh_out at once, so the SG_GRID_Y bands it holds are disjoint global - // row ranges; the whole workgroup then imageStores them. lr / MMA_M is the - // writing subgroup's warpInTile.y, so the global row reproduces the buffer - // path's gi exactly. - // - // PORTABILITY NOTE: the barrier() in the loop body keeps this loop rolled - // despite [[unroll]], so result[i][j] IS dynamically indexed. Coopmat arrays - // are opaque per-lane storage and dynamic indexing is exactly the construct - // the Xclipse/AMD-PAL compiler has broken before -- check this first if the - // texture variants miscompile on M51. - const uint CSH_TEXELS_PER_ROW = WG_TILE_N / 4u; - const uint CSH_TEXELS = CSH_ROWS * CSH_TEXELS_PER_ROW; - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - // Guards Csh_out against the previous iteration's readers. Inert on i == 0 - // but must stay unconditional to remain workgroup-uniform. - // coopmat-lds-fence: barrier() alone does NOT order shared stores against a - // subsequent coopMatLoad on the M51 Xclipse/AMD-PAL driver -- symptom is one - // stale MMA_M-row band of A, all columns, ~2.5% of runs, no crash (observed - // 2026-09-02 in sdpa_compute_out_coopmat.glsl). Measured cost of the fence: - // none (see this change's results). See memory - // `coopmat-lds-needs-explicit-memorybarriershared`. - memoryBarrierShared(); - barrier(); - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { -#ifdef HAS_BIAS - if (apply_bias > 0) { - const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopmat bias_tile; - coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); - result[i][j] += bias_tile; - } -#endif - coopmat out_tile = - coopmat(result[i][j]); - coopMatStore( - out_tile, Csh_out, - warpInTile.y * MMA_M * WG_TILE_N + - MMA_N * (MMAS_PER_SG_N * warpInTile.x + j), - WG_TILE_N, - gl_CooperativeMatrixLayoutRowMajor); - } - memoryBarrierShared(); - barrier(); - - for (uint t = gl_LocalInvocationID.x; t < CSH_TEXELS; t += WG_SIZE) { - const uint lr = t / CSH_TEXELS_PER_ROW; - const uint lc4 = t % CSH_TEXELS_PER_ROW; - const uint m = - tile_m_start + (lr / MMA_M) * SG_TILE_M + i * MMA_M + (lr % MMA_M); - const uint base = lr * WG_TILE_N + lc4 * 4u; - imageStore( - t_output, - ivec3(tile_n_start / 4u + lc4, m, 0), - vec4( - float(Csh_out[base]), - float(Csh_out[base + 1u]), - float(Csh_out[base + 2u]), - float(Csh_out[base + 3u]))); - } - } -#else - const uint N_out = uint(out_N_arg); - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint gi = tile_m_start + MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); - const uint gj = tile_n_start + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - -#ifdef HAS_BIAS - if (apply_bias > 0) { - const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopmat bias_tile; - coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); - result[i][j] += bias_tile; - } -#endif - - coopmat out_tile = - coopmat(result[i][j]); - coopMatStore( - out_tile, t_output, - gi * N_out + gj, N_out, - gl_CooperativeMatrixLayoutRowMajor); - } - } -#endif // IO_TEXTURE -} diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw.yaml b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw.yaml deleted file mode 100644 index ff97b1a2710..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw.yaml +++ /dev/null @@ -1,68 +0,0 @@ -# !!! KNOWN-INCORRECT -- DO NOT PROMOTE, DO NOT TIME !!! -# Fails the correctness gate deterministically (3/3 reps) on exactly the -# num_groups == 2 shapes (K=256: M128_K256_N128, M256_K256_N256). Every K=128, -# K=2048 and K=4096 case PASSES. Root cause NOT yet identified: the -# (slab, col) -> (n8blk, k4, component, parity) mapping was re-derived and is -# provably identical to the baseline slot map, so the defect is in something -# adjacent to the ownership change, not in the address arithmetic. -# Kept opt-in only (never a default) as a recorded negative result. -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -# "zpg" + "-tr" combination: tsweep_dbuf4zpg with its per-thread scalar -# A-staging replaced by tsweep_dbuf4tr's coopMat-mediated A staging. B -# staging, zp-hoist, byte-parallel nibble widening are all dbuf4zpg's, -# unchanged -- only A-staging differs. Requires t_packed_int8_input in the -# ROW-MAJOR kPackedInt8_4W layout (same requirement as tsweep_dbuf4tr). -# -# PROMOTED 2026-09-01 as the shipped default -- see QuantizedLinear.cpp's -# dq8ca_coopmat_variant() for the full validation record (10/10 buffer, 6/6 -# texture3d correctness across 1B/3B/8B; +4.2% real e2e prefill; 46.50% -> -# 49.94% efficiency on 8B). Also selectable explicitly via -# ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4zpgbw_txkgs. -# See openspec/changes/archive/2026-08-31-dq8ca-tr-staged-a-on-zpg. -# -# Single seed tile: the prior dbuf4zpg default's own tile (t128x64k32g42s32). -# A re-sweep against this shader's own (lower) register-pressure profile was -# run as a follow-up (coopmat-tr-tilesweep-4w-port) and found no better tile -# -- this remains the best known geometry. - -linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw: - parameter_names_with_default_values: - PRECISION: highp - HAS_BIAS: false - IO_STORAGE: buffer - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - MMA_M: 16 - MMA_N: 16 - MMA_K: 16 - WG_TILE_M: 64 - WG_TILE_N: 32 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - shader_variants: - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw_t128x64k32g42s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw_t128x64k32g42s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw2.glsl b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw2.glsl deleted file mode 100644 index 38640aca569..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw2.glsl +++ /dev/null @@ -1,762 +0,0 @@ -/* - * Copyright (c) Meta Platforms, Inc. and affiliates. - * All rights reserved. - * - * This source code is licensed under the BSD-style license found in the - * LICENSE file in the root directory of this source tree. - */ - -/* - * "zpg" + "-tr" combination: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpg.glsl - * with its per-thread scalar A-staging replaced by - * linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr.glsl's coopMat-mediated A - * staging (coopMatLoad(global) -> coopmat<> -> coopMatStore(LDS)). This is - * an ADDITIVE combination, not a redesign: every non-A-staging block below - * (B staging: coalesced write, no skew; zp-hoist: izp/ifs applied once after - * the group loop via wcorr_sh; byte-parallel nibble widening; static - * A_ALWAYS_ACTIVE branch elision -- N/A here, see below; group epilog; - * bias/store epilogue) is byte-identical to dbuf4zpg's. Only the A-staging - * block (prologue load+store, main-loop prefetch+store) is dbuf4tr's, - * verbatim. - * - * dbuf4zpg's per-thread A staging used an `a_active` guard (statically always - * true when A_ACTIVE_THREADS == WG_SIZE, via the A_MAP_FULL-gated - * A_ALWAYS_ACTIVE macro). dbuf4tr's per-SUBGROUP tile map has no equivalent - * concept -- every subgroup participates via a `t < NUM_A_TILES` guard that - * depends only on gl_SubgroupID, not gl_LocalInvocationID.x -- so - * A_MAP_FULL/A_ALWAYS_ACTIVE is dropped entirely in this file; it would be - * dead code for the new A-staging block. - * - * Rationale for combining this way (not the reverse) and why this is worth - * building at all: see this change's design.md D0-D3. In short -- the only - * existing measurement of dbuf4tr's A-staging technique (28.72-30.51%, - * dq8ca-arch-redesign) was taken against dbuf4tr's own pre-zpg baseline - * (old B skew, no byte-parallel widening, no branch elision) -- a materially - * weaker shader than the 46.49-46.50% dbuf4zpg this file now combines it - * with. This file exists to answer whether that combination performs - * differently now that register pressure is already reduced. - * - * A staging (the actual delta from dbuf4zpg): - * dbuf4zpg: per-thread (m4, k4) ivec4 fetch, hoisted a_lds_off0/a_glb_row; - * only A_ACTIVE_THREADS invocations participate, each scattering - * 4 rows into Ash_int8 with 4 scalar stores. - * this file: per-SUBGROUP MMA_M x MMA_K tile fetch via coopMatLoad straight - * from a ROW-MAJOR (kPackedInt8_4W) int8 activation buffer, then - * coopMatStore into the same Ash_int8 slot -- dbuf4tr's mapping, - * unmodified (not re-derived; see design.md D3). - * - * t_packed_int8_input is therefore bound the same way dbuf4tr binds it: a - * SCALAR int8_t array in the kPackedInt8_4W layout (plain row-major int8, - * row stride K), produced by quantize_and_pack_4w_with_group_sums.glsl. - * QuantizedLinear.cpp's dq8ca_variant_wants_rowmajor_a() must recognize this - * file's variant token (tsweep_dbuf4zpgtr_t...) the same way it already - * recognizes tsweep_dbuf4tr_t/trm_t/trd_t, so graph-build time (packer - * selection) and dispatch time (kernel selection) cannot disagree. - * - * B CANNOT be coopmat-staged (int4 nibble unpack; a coopmat's per-lane layout - * is opaque to hand-assembly from unpacked registers) -- unchanged from both - * parent files. B staging below is dbuf4zpg's byte-parallel, coalesced, - * no-skew version, untouched. - * - * The loop structure is dbuf4's (both parents share it), unchanged: - * prologue: prefetch chunk 0 -> temp, store to slice 0 (no barrier) - * per iter: barrier -> prefetch(next) -> MMA(cur) -> store(next) - * kept nested (groups x chunks) with an unconditional group epilog -- - * flattening it crashes the Xclipse PAL compiler at large spec-resolved trip - * counts (see dbuf2's own header). - * - * Selected via - * ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4zpgtr_txkgs<32|64> - * (QuantizedLinear.cpp), additive to the tsweep_dbuf4zpg_t..., tsweep_dbuf4tr_t... - * and tsweep_t... namespaces. NOT the default -- unvalidated until it passes - * repeated test_llama_microbench --correctness-only runs (see - * dq8ca_coopmat_variant()'s comment on why a single pass is not proof). - * - * Performs: out[M,N] = dequant(int8_act) * dequant(int_w) (+ bias) - * via coopmat x coopmat -> coopmat on the matrix unit. - * - * Hard preconditions (dbuf4zpg's, plus dbuf4tr's row-major/alignment ones): - * M % WG_TILE_M == 0, N % WG_TILE_N == 0, K % WG_TILE_K == 0, - * group_size % WG_TILE_K == 0, K % 4 == 0, - * WG_TILE_M % MMA_M == 0, WG_TILE_K % MMA_K == 0, - * t_packed_int8_input in kPackedInt8_4W (row-major) layout, - * device exposes coopmatx-> at 16x16x16. - */ - -#version 450 core - -#extension GL_KHR_cooperative_matrix : require -#extension GL_KHR_memory_scope_semantics : require -#extension GL_KHR_shader_subgroup_basic : enable -#extension GL_EXT_shader_explicit_arithmetic_types : require -#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require -// 8-bit SSBO access: A is bound as a scalar int8_t array so that the -// coopMatLoad below has a MATCHING component type (see dbuf4tr's header for -// why the type must match on this driver). -#extension GL_EXT_shader_8bit_storage : require -#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require -#extension GL_EXT_control_flow_attributes : enable - -#define PRECISION ${PRECISION} - -$if WEIGHT_NBITS == 4: - #define WEIGHT_INT4 - -$if HAS_BIAS: - #define HAS_BIAS - -$if WEIGHT_STORAGE == "buffer": - #define WEIGHT_BUFFER - -$if IO_STORAGE == "texture3d": - #define IO_TEXTURE - -layout(std430) buffer; - -#include "common.glslh" - -// Bindings — match add_linear_dqa_qw_node arg order: -// output(0), fp_input(1), packed_int8_input(2), int_input_sums(3 - unused), -// input_scales(4), input_zps(5), packed_weight(6), weight_sums(7), -// weight_scales(8), bias(9). -${layout_declare_tensor(B, "w", "t_output", "half", IO_STORAGE, is_scalar_array=True)} -// t_input is unread here -- the activations arrive already quantized in -// t_packed_int8_input -- but stays declared so the binding layout matches the -// dispatch site. It tracks IO_STORAGE so the two IO tensors stay consistent. -${layout_declare_tensor(B, "r", "t_input", "half", IO_STORAGE, is_scalar_array=False)} -// ROW-MAJOR (kPackedInt8_4W) packed activations, bound as a scalar int8_t -// array (row stride = K int8) -- dbuf4tr's binding, unchanged. The stock -// 4h4w layout dbuf4zpg uses is NOT row-major (component index selects a row, -// non-affine), so it cannot be addressed by any coopMatLoad. -${layout_declare_tensor(B, "r", "t_packed_int8_input", "int8", "buffer", is_scalar_array=True)} -${layout_declare_tensor(B, "r", "t_int8_input_sums", "int", "buffer", is_scalar_array=True)} -${layout_declare_tensor(B, "r", "t_int8_input_scales", "half", "texture3d")} -${layout_declare_tensor(B, "r", "t_int8_input_zps", "int8", "texture3d")} -${layout_declare_tensor(B, "r", "t_packed_weight", "int", WEIGHT_STORAGE, is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_weight_sums", "int", "buffer", is_scalar_array=True)} -${layout_declare_tensor(B, "r", "t_weight_scales", "half", "buffer", is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_bias", "half", "buffer", is_scalar_array=True)} - -${layout_declare_ubo(B, "ivec4", "output_sizes")} -${layout_declare_ubo(B, "ivec4", "input_sizes")} - -layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; - -${layout_declare_spec_const(C, "int", "apply_bias", "0")} -// INT4 only; inert (0) for INT8 so the dispatcher's spec list lines up. -${layout_declare_spec_const(C, "int", "K4_per_group", "0")} -${layout_declare_spec_const(C, "int", "num_groups_arg", "0")} -${layout_declare_spec_const(C, "int", "out_N_arg", "0")} - -// Tile geometry -const uint MMA_M = ${MMA_M}; -const uint MMA_N = ${MMA_N}; -const uint MMA_K = ${MMA_K}; - -const uint WG_TILE_M = ${WG_TILE_M}; -const uint WG_TILE_N = ${WG_TILE_N}; -const uint WG_TILE_K = ${WG_TILE_K}; - -const uint SG_GRID_X = ${SG_GRID_X}; -const uint SG_GRID_Y = ${SG_GRID_Y}; -const uint SUBGROUP_SIZE = ${SUBGROUP_SIZE}; -const uint NUM_SUBGROUPS = SG_GRID_X * SG_GRID_Y; -const uint WG_SIZE = NUM_SUBGROUPS * SUBGROUP_SIZE; - -const uint SG_TILE_M = WG_TILE_M / SG_GRID_Y; -const uint SG_TILE_N = WG_TILE_N / SG_GRID_X; -const uint MMAS_PER_SG_M = SG_TILE_M / MMA_M; -const uint MMAS_PER_SG_N = SG_TILE_N / MMA_N; - -const uint A_SLAB_INT8 = WG_TILE_M * MMA_K; -const uint B_USEFUL_U32 = MMA_K / 4u; -// No skew + coalesced write -- dbuf4zpg's B fix, unchanged (this file does -// not touch B staging at all). -const uint B_STRIDE_U32 = B_USEFUL_U32; -const uint B_SLAB_U32 = WG_TILE_N * B_STRIDE_U32; -const uint NUM_K_SLABS = WG_TILE_K / MMA_K; - -const uint A_SLAB_U32 = A_SLAB_INT8 / 4u; -const uint A_STRIDE_U32 = MMA_K / 4u; - -// One ping-pong slice covers all K-slabs of one chunk. -const uint ASH_SLICE_U32 = NUM_K_SLABS * A_SLAB_U32; -const uint BSH_SLICE_U32 = NUM_K_SLABS * B_SLAB_U32; - -// Double-buffered MMA operand staging. -shared uint Ash_int8[2u * ASH_SLICE_U32]; -shared uint Bsh_int8[2u * BSH_SLICE_U32]; - -// Per-WG-tile-row activation params (loaded ONCE at WG start; constant -// across groups). -shared int izp_sh[WG_TILE_M]; // int32 (cast from int8 source) for broadcast -shared float ifs_sh[WG_TILE_M]; // float32 (cast from fp16 source) for broadcast - -// Per-(group, output-channel) weight params, ping-ponged by group parity. -// (For per-channel INT8 only slice 0 is ever used.) -shared float wsc_sh[2u * WG_TILE_N]; -// SUM_g wsc[g][n]*wsum[g][n] per output channel -- weight-side only, so it is -// accumulated once in the prologue. dbuf4zpg's zp-hoist, unchanged. -shared float wcorr_sh[WG_TILE_N]; - -#ifdef HAS_BIAS -shared float bias_sh[WG_TILE_N]; -#endif - -#ifdef IO_TEXTURE -// Result staging for the imageStore epilogue, mirroring the fp16 kernel: -// SG_GRID_Y bands of MMA_M rows, each WG_TILE_N wide, row-major. A full -// WG_TILE_M x WG_TILE_N buffer would cost SG_GRID_Y/MMAS_PER_SG_M x more LDS -// and wreck occupancy. float16_t-typed because coopMatStore needs it. -const uint CSH_ROWS = SG_GRID_Y * MMA_M; -shared float16_t Csh_out[CSH_ROWS * WG_TILE_N]; -#endif - -// Running fp32 accumulator (across all groups). -coopmat - result[MMAS_PER_SG_M][MMAS_PER_SG_N]; - -// Per-group int32 MMA accumulator. -coopmat - accum_int32[MMAS_PER_SG_M][MMAS_PER_SG_N]; - - -// Byte-parallel int4 -> int8 widening. dbuf4zpg's, unchanged (B-side only). -// -// The four nibbles this shader needs from one packed uint are ALREADY one per -// byte (bits 3:0 of each byte for parity 0, bits 7:4 for parity 1), so all four -// can be widened at once instead of with a per-nibble -// shift/mask/bias-subtract/mask chain. -// -// For v in [0,15] the biased value is v-8. `v ^ 8` is exactly the 4-bit -// two's-complement pattern of v-8, because -8 == +8 (mod 16): -// v=0 -> 0x8 -> -8 v=7 -> 0xF -> -1 -// v=8 -> 0x0 -> 0 v=15 -> 0x7 -> +7 -// so the only remaining work is sign-extending bit 3 into bits 7:4 per byte. -// `sgn * 0x1E` does that with no cross-byte carry: 0x08 * 0x1E == 0xF0 exactly, -// and sgn is at most 0x08080808 so the product is at most 0xF0F0F0F0. -// -// A naive `nib - 0x08080808` would NOT work -- it borrows across byte lanes -// whenever a nibble is < 8. Shifts must be on uint, not int, so the >> is -// logical rather than arithmetic. -// -// ~5 ops per 4 weights vs ~22 for the per-nibble chain; bit-identical output. -uint widen_nibbles(const uint w, const uint parity) { - const uint nib = (parity == 0u) ? (w & 0x0F0F0F0Fu) : ((w >> 4u) & 0x0F0F0F0Fu); - const uint p = nib ^ 0x08080808u; - const uint sgn = p & 0x08080808u; - return p | (sgn * 0x1Eu); -} - -void main() { - const uvec2 tileID = uvec2(gl_WorkGroupID.xy); - const uvec2 warpInTile = uvec2( - gl_SubgroupID % SG_GRID_X, - gl_SubgroupID / SG_GRID_X); - - const uint K = uint(input_sizes.x); - const uint N = uint(output_sizes.x); - const uint N4 = (N + 3u) / 4u; - const uint nblocks_x_A = (K + 3u) >> 2u; - // A row stride in INT8 elements (dbuf4tr's binding is row-major int8, not - // the 4h4w ivec4 block layout dbuf4zpg used -- so A addressing below is in - // int8 elements, not int, and derived from nblocks_x_A so it matches the - // packer's `m_row * K4 + k4` addressing exactly (K % 4 == 0 makes them - // equal to K directly). - const uint a_row_stride_i8 = nblocks_x_A * 4u; - -#ifdef WEIGHT_INT4 - const uint num_groups = uint(num_groups_arg); - const uint CHUNKS_PER_GROUP = uint(K4_per_group) * 4u / WG_TILE_K; -#else - // Per-channel: a single quant "group" spanning all of K. The nested - // groups x chunks loop below collapses to a flat chunk loop, the wsum/wsc - // ping-pong never crosses a boundary, and the epilog runs exactly once. - const uint num_groups = 1u; - const uint CHUNKS_PER_GROUP = uint(num_groups_arg); -#endif - const uint num_chunks = num_groups * CHUNKS_PER_GROUP; - - const uint tile_m_start = WG_TILE_M * tileID.y; - const uint tile_n_start = WG_TILE_N * tileID.x; - - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - result[i][j] = coopmat(0.0); - accum_int32[i][j] = coopmat(0); - } - } - - const uint K_BLOCKS_PER_CHUNK = WG_TILE_K >> 2u; - - // --- A staging tile map (dbuf4tr's, unmodified): one MMA_M x MMA_K coopmat - // tile per subgroup per slot, dealt round-robin across the - // NUM_SUBGROUPS subgroups so every subgroup participates. Replaces - // dbuf4zpg's per-thread (m4, k4) map / a_active guard entirely -- see - // design.md D3 for why this is reused as-is, not re-derived. - const uint A_TILES_M = WG_TILE_M / MMA_M; - const uint A_TILES_K = WG_TILE_K / MMA_K; // == NUM_K_SLABS - const uint NUM_A_TILES = A_TILES_M * A_TILES_K; - const uint A_TILES_PER_SG = (NUM_A_TILES + NUM_SUBGROUPS - 1u) / NUM_SUBGROUPS; - -#ifdef WEIGHT_INT4 - // --- B staging thread map: (block, col) slots; each slot extracts one - // ColumnMajor LDS uint (4 K-contiguous sign-extended int8) --- - const uint B_TOTAL_SLOTS = K_BLOCKS_PER_CHUNK * WG_TILE_N; - const uint B_SLOTS_PER_THREAD = B_TOTAL_SLOTS / WG_SIZE; - const uint N8_PER_TILE = WG_TILE_N >> 3u; -#else - // --- B staging thread map: one (k4, n4) ivec4 block per active thread --- - const uint B_FETCH_SLOTS = K_BLOCKS_PER_CHUNK * (WG_TILE_N >> 2u); - const uint N4_PER_TILE = WG_TILE_N >> 2u; - const uint b_k4_in_chunk = gl_LocalInvocationID.x / N4_PER_TILE; - const uint b_n_uint_col = gl_LocalInvocationID.x % N4_PER_TILE; - const bool b_active = gl_LocalInvocationID.x < B_FETCH_SLOTS; -#endif - - // ===== INTERVENTION F: hoist loop-invariant B staging index math ===== - // dbuf4zpg's, unchanged -- B staging is untouched by this file's A-staging - // swap. See dbuf4zpg's header for the full rationale (ablation-attributed - // -16.8% block, mostly index arithmetic not memory traffic). -#ifdef WEIGHT_INT4 - // ===== bw2 DIAGNOSTIC: 2 consecutive uints per thread, ALL 256 threads - // active (no idling). Discriminates "4-wide ownership + idle threads" from - // "consecutive-write ownership per se" as the cause of the num_groups==2 - // failure seen in bw/bwr. - const uint B_PAIR = 2u; - const uint b_lin0 = gl_LocalInvocationID.x * B_PAIR; // first uint index - const uint b_slab = b_lin0 / B_SLAB_U32; - const uint b_inslab = b_lin0 % B_SLAB_U32; - const uint b_col = b_inslab / B_STRIDE_U32; - const uint b_k4lo = b_inslab % B_STRIDE_U32; - const uint b_rem = b_col & 7u; - const uint b_comp1 = b_rem & 3u; - const uint b_par1 = b_rem >> 2u; - const uint b_n8blk1 = (tile_n_start >> 3u) + (b_col >> 3u); - const uint b_k4base = b_slab * (MMA_K >> 2u) + b_k4lo; - const uint b_lds_base = b_lin0; -#endif - - // Prefetch temp registers. temp_A is a coopmat array (dbuf4tr's A-staging - // technique); indices into it are [[unroll]]-resolved compile-time - // constants, never dynamic -- dynamic indexing of a coopmat array is - // exactly the construct the Xclipse/AMD-PAL compiler has miscompiled - // before. - coopmat - temp_A[A_TILES_PER_SG]; -#ifdef WEIGHT_INT4 - ivec4 temp_B[2u]; // bw2: one per k4 of the owned pair - float temp_wsc; -#else - ivec4 temp_B; -#endif - - // ========================================================= - // PROLOGUE - // ========================================================= - if (gl_LocalInvocationID.x < (WG_TILE_M >> 2u)) { - const uint m4 = (tile_m_start >> 2u) + gl_LocalInvocationID.x; - const vec4 sc = vec4(texelFetch(t_int8_input_scales, ivec3(m4, 0, 0), 0)); - const ivec4 zp = texelFetch(t_int8_input_zps, ivec3(m4, 0, 0), 0); - const uint base = gl_LocalInvocationID.x * 4u; - ifs_sh[base + 0u] = sc.x; ifs_sh[base + 1u] = sc.y; - ifs_sh[base + 2u] = sc.z; ifs_sh[base + 3u] = sc.w; - izp_sh[base + 0u] = zp.x; izp_sh[base + 1u] = zp.y; - izp_sh[base + 2u] = zp.z; izp_sh[base + 3u] = zp.w; - } - // Group 0 weight scales -> slice 0, and the hoisted weight-side correction - // SUM_g wsc[g][n]*wsum[g][n] accumulated across ALL groups. dbuf4zpg's - // zp-hoist, unchanged. - if (gl_LocalInvocationID.x < WG_TILE_N) { - const uint n_idx = tile_n_start + gl_LocalInvocationID.x; - f16vec4 sv0 = t_weight_scales[n_idx >> 2u]; - wsc_sh[gl_LocalInvocationID.x] = float(sv0[n_idx & 3u]); - - float corr = 0.0; - for (uint g = 0; g < num_groups; ++g) { - f16vec4 sv = t_weight_scales[g * N4 + (n_idx >> 2u)]; - corr += float(sv[n_idx & 3u]) * float(t_weight_sums[g * N + n_idx]); - } - wcorr_sh[gl_LocalInvocationID.x] = corr; - } - memoryBarrierShared(); - barrier(); - - // NOTE: dbuf4zpg builds izp_bcast/ifs_bcast AFTER the group loop, not here - // -- that is the register-pressure saving zp-hoist buys. Unchanged. - - // dbuf4: prefetch chunk 0 into temp registers, THEN store to slice 0 (no - // barrier here -- the main loop's first iteration barriers before - // reading slice 0). - // - // A staging (dbuf4tr's technique): per-subgroup coopMatLoad straight from - // the row-major global buffer. - [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { - const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; - if (t < NUM_A_TILES) { - const uint tm = t / A_TILES_K; - const uint tk = t % A_TILES_K; - coopMatLoad( - temp_A[s], t_packed_int8_input, - (tile_m_start + tm * MMA_M) * a_row_stride_i8 + tk * MMA_K, - a_row_stride_i8, - gl_CooperativeMatrixLayoutRowMajor); - } - } -#ifdef WEIGHT_INT4 - [[unroll]] for (uint kk = 0; kk < 2u; ++kk) { -#ifdef WEIGHT_BUFFER - temp_B[kk] = t_packed_weight[(b_n8blk1 * nblocks_x_A) + b_k4base + kk]; -#else - temp_B[kk] = texelFetch(t_packed_weight, ivec2(b_k4base + kk, b_n8blk1), 0); -#endif - } -#else - if (b_active) { - const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; -#ifdef WEIGHT_BUFFER - temp_B = t_packed_weight[(b_k4_in_chunk * N4) + block_x_w]; -#else - temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, b_k4_in_chunk), 0); -#endif - } -#endif - { - // store chunk 0 -> slice 0 - // A staging (dbuf4tr's technique): coopMatStore into the same Ash_int8 - // slot layout dbuf4zpg's scalar scatter used to write. - [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { - const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; - if (t < NUM_A_TILES) { - const uint tm = t / A_TILES_K; - const uint tk = t % A_TILES_K; - coopMatStore( - temp_A[s], Ash_int8, - tk * A_SLAB_U32 + (tm * MMA_M) * A_STRIDE_U32, - A_STRIDE_U32, - gl_CooperativeMatrixLayoutRowMajor); - } - } -#ifdef WEIGHT_INT4 - [[unroll]] for (uint kk = 0; kk < 2u; ++kk) { - Bsh_int8[b_lds_base + kk] = - widen_nibbles(uint(temp_B[kk][b_comp1]), b_par1); - } -#else - if (b_active) { - const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); - const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); - const uint n_col_base = b_n_uint_col * 4u; - [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { - Bsh_int8[slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab] = - uint(temp_B[n_in_blk]); - } - } -#endif - } - - // ========================================================= - // MAIN LOOP (dbuf4) — nested groups x chunks (kept nested; flattening it - // with a conditional coopmat epilog crashes the Xclipse PAL compiler at - // large spec-resolved trip counts). One barrier per chunk. Chunk - // iteration (global index `chunk`): - // 1. barrier — A/B slice (chunk%2) fully written; on the first chunk - // of group g, wsc slice (g%2) is too. - // 2. prefetch — chunk+1 (A tiles, B blocks) into temp; when chunk+1 - // starts a new group, also its wsc element. Skipped - // entirely on the final chunk. - // 3. int8 MMA — on slice (chunk%2) into accum_int32. - // 4. store — temp -> A/B slice ((chunk+1)%2), unpacking the weight; - // on a group boundary, wsc -> slice ((g+1)%2). - // The group epilog runs unconditionally at the tail of each group. - // ========================================================= - uint chunk = 0; - for (uint group_i = 0; group_i < num_groups; ++group_i) { - for (uint inner = 0; inner < CHUNKS_PER_GROUP; ++inner, ++chunk) { - const bool has_next = chunk + 1u < num_chunks; - const bool group_crossing = has_next && (inner + 1u == CHUNKS_PER_GROUP); - const uint cur_a = (chunk % 2u) * ASH_SLICE_U32; - const uint cur_b = (chunk % 2u) * BSH_SLICE_U32; - const uint nxt_a = ((chunk + 1u) % 2u) * ASH_SLICE_U32; - const uint nxt_b = ((chunk + 1u) % 2u) * BSH_SLICE_U32; - - // coopmat-lds-fence: barrier() alone does NOT order shared stores against a - // subsequent coopMatLoad on the M51 Xclipse/AMD-PAL driver -- symptom is one - // stale MMA_M-row band of A, all columns, ~2.5% of runs, no crash (observed - // 2026-09-02 in sdpa_compute_out_coopmat.glsl). Measured cost of the fence: - // none (see this change's results). See memory - // `coopmat-lds-needs-explicit-memorybarriershared`. - memoryBarrierShared(); - barrier(); - - // --- 2. prefetch chunk+1 -> temp --- - if (has_next) { - const uint chunkK_nxt = (chunk + 1u) * WG_TILE_K; - // A staging (dbuf4tr's technique): coopMatLoad straight from global. - [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { - const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; - if (t < NUM_A_TILES) { - const uint tm = t / A_TILES_K; - const uint tk = t % A_TILES_K; - coopMatLoad( - temp_A[s], t_packed_int8_input, - (tile_m_start + tm * MMA_M) * a_row_stride_i8 + chunkK_nxt + - tk * MMA_K, - a_row_stride_i8, - gl_CooperativeMatrixLayoutRowMajor); - } - } -#ifdef WEIGHT_INT4 - [[unroll]] for (uint kk = 0; kk < 2u; ++kk) { - const uint k4_blk = (chunkK_nxt >> 2u) + b_k4base + kk; -#ifdef WEIGHT_BUFFER - temp_B[kk] = t_packed_weight[(b_n8blk1 * nblocks_x_A) + k4_blk]; -#else - temp_B[kk] = texelFetch(t_packed_weight, ivec2(k4_blk, b_n8blk1), 0); -#endif - } - if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { - const uint n_idx = tile_n_start + gl_LocalInvocationID.x; - f16vec4 sv = t_weight_scales[(group_i + 1u) * N4 + (n_idx >> 2u)]; - temp_wsc = float(sv[n_idx & 3u]); - } -#else - if (b_active) { - const uint block_y_w = (chunkK_nxt >> 2u) + b_k4_in_chunk; - const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; -#ifdef WEIGHT_BUFFER - temp_B = t_packed_weight[(block_y_w * N4) + block_x_w]; -#else - temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, block_y_w), 0); -#endif - } -#endif - } - - // --- 3. int8 MMA on the cur slice --- - [[unroll]] for (uint k = 0; k < NUM_K_SLABS; ++k) { - const uint slab_a_base_u32 = cur_a + k * A_SLAB_U32; - const uint slab_b_base_u32 = cur_b + k * B_SLAB_U32; - - coopmat matA[MMAS_PER_SG_M]; - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - const uint row_a = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); - coopMatLoad( - matA[i], Ash_int8, - slab_a_base_u32 + row_a * A_STRIDE_U32, - A_STRIDE_U32, - gl_CooperativeMatrixLayoutRowMajor); - } - - coopmat matB; - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint col_b = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopMatLoad( - matB, Bsh_int8, - slab_b_base_u32 + col_b * B_STRIDE_U32, - B_STRIDE_U32, - gl_CooperativeMatrixLayoutColumnMajor); - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - accum_int32[i][j] = coopMatMulAdd(matA[i], matB, accum_int32[i][j]); - } - } - } - - // --- 4. store temp (chunk+1) -> nxt slice --- - if (has_next) { - // A staging (dbuf4tr's technique): coopMatStore into the nxt slice. - [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { - const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; - if (t < NUM_A_TILES) { - const uint tm = t / A_TILES_K; - const uint tk = t % A_TILES_K; - coopMatStore( - temp_A[s], Ash_int8, - nxt_a + tk * A_SLAB_U32 + (tm * MMA_M) * A_STRIDE_U32, - A_STRIDE_U32, - gl_CooperativeMatrixLayoutRowMajor); - } - } -#ifdef WEIGHT_INT4 - [[unroll]] for (uint kk = 0; kk < 2u; ++kk) { - Bsh_int8[nxt_b + b_lds_base + kk] = - widen_nibbles(uint(temp_B[kk][b_comp1]), b_par1); - } - if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { - const uint wbase_nxt = ((group_i + 1u) % 2u) * WG_TILE_N; - wsc_sh[wbase_nxt + gl_LocalInvocationID.x] = temp_wsc; - } -#else - if (b_active) { - const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); - const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); - const uint n_col_base = b_n_uint_col * 4u; - [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { - Bsh_int8[nxt_b + slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab] = - uint(temp_B[n_in_blk]); - } - } -#endif - } - } // chunks - - // --- Group epilog: scale-only accumulate, reset accum --- - // dbuf4zpg's, unchanged. Just result += float(acc) * wsc. The - // zero-point subtract and the ifs multiply are hoisted out of the group - // loop (applied once below). - { - const uint wbase = (group_i % 2u) * WG_TILE_N; - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - - coopmat wsc_bcast; - coopMatLoad( - wsc_bcast, wsc_sh, - wbase + local_n_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutRowMajor); - - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - result[i][j] += - coopmat( - accum_int32[i][j]) * wsc_bcast; - accum_int32[i][j] = coopmat(0); - } - } - } - } // groups - - // --- Hoisted correction, applied ONCE: -------------------------------- - // result = ifs * ( result - izp * SUM_g wsc_g*wsum_g ) - // dbuf4zpg's, unchanged. izp/ifs are loaded here rather than before the - // group loop so they are not live across it. - { - coopmat - izpf_bcast[MMAS_PER_SG_M]; - coopmat - ifs_bcast[MMAS_PER_SG_M]; - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - const uint local_m_base = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); - coopmat izp_i; - coopMatLoad( - izp_i, izp_sh, local_m_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutColumnMajor); - izpf_bcast[i] = - coopmat(izp_i); - coopMatLoad( - ifs_bcast[i], ifs_sh, local_m_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutColumnMajor); - } - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopmat wcorr_bcast; - coopMatLoad( - wcorr_bcast, wcorr_sh, local_n_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutRowMajor); - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - result[i][j] = ifs_bcast[i] * (result[i][j] - izpf_bcast[i] * wcorr_bcast); - } - } - } - - // --- Bias (optional) --- -#ifdef HAS_BIAS - if (apply_bias > 0) { - for (uint t = gl_LocalInvocationID.x; t < WG_TILE_N; t += WG_SIZE) { - bias_sh[t] = float(t_bias[tile_n_start + t]); - } - memoryBarrierShared(); - barrier(); - } -#endif - - // --- Store result tile --- - // N for the store address math MUST come from the spec constant, not the - // sizes UBO (see out_N_arg above). -#ifdef IO_TEXTURE - // Epilogue iteration i drains accumulator row-block i from EVERY subgroup - // into Csh_out at once, so the SG_GRID_Y bands it holds are disjoint global - // row ranges; the whole workgroup then imageStores them. lr / MMA_M is the - // writing subgroup's warpInTile.y, so the global row reproduces the buffer - // path's gi exactly. - // - // PORTABILITY NOTE: the barrier() in the loop body keeps this loop rolled - // despite [[unroll]], so result[i][j] IS dynamically indexed. Coopmat arrays - // are opaque per-lane storage and dynamic indexing is exactly the construct - // the Xclipse/AMD-PAL compiler has broken before -- check this first if the - // texture variants miscompile on M51. - const uint CSH_TEXELS_PER_ROW = WG_TILE_N / 4u; - const uint CSH_TEXELS = CSH_ROWS * CSH_TEXELS_PER_ROW; - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - // Guards Csh_out against the previous iteration's readers. Inert on i == 0 - // but must stay unconditional to remain workgroup-uniform. - // coopmat-lds-fence: barrier() alone does NOT order shared stores against a - // subsequent coopMatLoad on the M51 Xclipse/AMD-PAL driver -- symptom is one - // stale MMA_M-row band of A, all columns, ~2.5% of runs, no crash (observed - // 2026-09-02 in sdpa_compute_out_coopmat.glsl). Measured cost of the fence: - // none (see this change's results). See memory - // `coopmat-lds-needs-explicit-memorybarriershared`. - memoryBarrierShared(); - barrier(); - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { -#ifdef HAS_BIAS - if (apply_bias > 0) { - const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopmat bias_tile; - coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); - result[i][j] += bias_tile; - } -#endif - coopmat out_tile = - coopmat(result[i][j]); - coopMatStore( - out_tile, Csh_out, - warpInTile.y * MMA_M * WG_TILE_N + - MMA_N * (MMAS_PER_SG_N * warpInTile.x + j), - WG_TILE_N, - gl_CooperativeMatrixLayoutRowMajor); - } - memoryBarrierShared(); - barrier(); - - for (uint t = gl_LocalInvocationID.x; t < CSH_TEXELS; t += WG_SIZE) { - const uint lr = t / CSH_TEXELS_PER_ROW; - const uint lc4 = t % CSH_TEXELS_PER_ROW; - const uint m = - tile_m_start + (lr / MMA_M) * SG_TILE_M + i * MMA_M + (lr % MMA_M); - const uint base = lr * WG_TILE_N + lc4 * 4u; - imageStore( - t_output, - ivec3(tile_n_start / 4u + lc4, m, 0), - vec4( - float(Csh_out[base]), - float(Csh_out[base + 1u]), - float(Csh_out[base + 2u]), - float(Csh_out[base + 3u]))); - } - } -#else - const uint N_out = uint(out_N_arg); - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint gi = tile_m_start + MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); - const uint gj = tile_n_start + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - -#ifdef HAS_BIAS - if (apply_bias > 0) { - const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopmat bias_tile; - coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); - result[i][j] += bias_tile; - } -#endif - - coopmat out_tile = - coopmat(result[i][j]); - coopMatStore( - out_tile, t_output, - gi * N_out + gj, N_out, - gl_CooperativeMatrixLayoutRowMajor); - } - } -#endif // IO_TEXTURE -} diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw2.yaml b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw2.yaml deleted file mode 100644 index c51f9cb0359..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw2.yaml +++ /dev/null @@ -1,60 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -# "zpg" + "-tr" combination: tsweep_dbuf4zpg with its per-thread scalar -# A-staging replaced by tsweep_dbuf4tr's coopMat-mediated A staging. B -# staging, zp-hoist, byte-parallel nibble widening are all dbuf4zpg's, -# unchanged -- only A-staging differs. Requires t_packed_int8_input in the -# ROW-MAJOR kPackedInt8_4W layout (same requirement as tsweep_dbuf4tr). -# -# PROMOTED 2026-09-01 as the shipped default -- see QuantizedLinear.cpp's -# dq8ca_coopmat_variant() for the full validation record (10/10 buffer, 6/6 -# texture3d correctness across 1B/3B/8B; +4.2% real e2e prefill; 46.50% -> -# 49.94% efficiency on 8B). Also selectable explicitly via -# ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4zpgbw2_txkgs. -# See openspec/changes/archive/2026-08-31-dq8ca-tr-staged-a-on-zpg. -# -# Single seed tile: the prior dbuf4zpg default's own tile (t128x64k32g42s32). -# A re-sweep against this shader's own (lower) register-pressure profile was -# run as a follow-up (coopmat-tr-tilesweep-4w-port) and found no better tile -# -- this remains the best known geometry. - -linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw2: - parameter_names_with_default_values: - PRECISION: highp - HAS_BIAS: false - IO_STORAGE: buffer - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - MMA_M: 16 - MMA_N: 16 - MMA_K: 16 - WG_TILE_M: 64 - WG_TILE_N: 32 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - shader_variants: - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw2_t128x64k32g42s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw2_t128x64k32g42s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw3.glsl b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw3.glsl deleted file mode 100644 index 5f71b6ed5e1..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw3.glsl +++ /dev/null @@ -1,798 +0,0 @@ -/* - * Copyright (c) Meta Platforms, Inc. and affiliates. - * All rights reserved. - * - * This source code is licensed under the BSD-style license found in the - * LICENSE file in the root directory of this source tree. - */ - -/* - * "zpg" + "-tr" combination: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpg.glsl - * with its per-thread scalar A-staging replaced by - * linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr.glsl's coopMat-mediated A - * staging, PLUS "bw3": shared B retyped to uvec2 with dbuf4zpgbw2's - * all-threads-active ownership, so each thread issues ONE ds_write_b64 for its - * 2-uint pair. A staging, the quantization math and every safeguard are - * dbuf4zpgtr's, unchanged. This is - * an ADDITIVE combination, not a redesign: every non-A-staging block below - * (B staging: coalesced write, no skew; zp-hoist: izp/ifs applied once after - * the group loop via wcorr_sh; byte-parallel nibble widening; static - * A_ALWAYS_ACTIVE branch elision -- N/A here, see below; group epilog; - * bias/store epilogue) is byte-identical to dbuf4zpg's. Only the A-staging - * block (prologue load+store, main-loop prefetch+store) is dbuf4tr's, - * verbatim. - * - * dbuf4zpg's per-thread A staging used an `a_active` guard (statically always - * true when A_ACTIVE_THREADS == WG_SIZE, via the A_MAP_FULL-gated - * A_ALWAYS_ACTIVE macro). dbuf4tr's per-SUBGROUP tile map has no equivalent - * concept -- every subgroup participates via a `t < NUM_A_TILES` guard that - * depends only on gl_SubgroupID, not gl_LocalInvocationID.x -- so - * A_MAP_FULL/A_ALWAYS_ACTIVE is dropped entirely in this file; it would be - * dead code for the new A-staging block. - * - * Rationale for combining this way (not the reverse) and why this is worth - * building at all: see this change's design.md D0-D3. In short -- the only - * existing measurement of dbuf4tr's A-staging technique (28.72-30.51%, - * dq8ca-arch-redesign) was taken against dbuf4tr's own pre-zpg baseline - * (old B skew, no byte-parallel widening, no branch elision) -- a materially - * weaker shader than the 46.49-46.50% dbuf4zpg this file now combines it - * with. This file exists to answer whether that combination performs - * differently now that register pressure is already reduced. - * - * A staging (the actual delta from dbuf4zpg): - * dbuf4zpg: per-thread (m4, k4) ivec4 fetch, hoisted a_lds_off0/a_glb_row; - * only A_ACTIVE_THREADS invocations participate, each scattering - * 4 rows into Ash_int8 with 4 scalar stores. - * this file: per-SUBGROUP MMA_M x MMA_K tile fetch via coopMatLoad straight - * from a ROW-MAJOR (kPackedInt8_4W) int8 activation buffer, then - * coopMatStore into the same Ash_int8 slot -- dbuf4tr's mapping, - * unmodified (not re-derived; see design.md D3). - * - * t_packed_int8_input is therefore bound the same way dbuf4tr binds it: a - * SCALAR int8_t array in the kPackedInt8_4W layout (plain row-major int8, - * row stride K), produced by quantize_and_pack_4w_with_group_sums.glsl. - * QuantizedLinear.cpp's dq8ca_variant_wants_rowmajor_a() must recognize this - * file's variant token (tsweep_dbuf4zpgtr_t...) the same way it already - * recognizes tsweep_dbuf4tr_t/trm_t/trd_t, so graph-build time (packer - * selection) and dispatch time (kernel selection) cannot disagree. - * - * B CANNOT be coopmat-staged (int4 nibble unpack; a coopmat's per-lane layout - * is opaque to hand-assembly from unpacked registers) -- unchanged from both - * parent files. B staging below is dbuf4zpg's byte-parallel, coalesced, - * no-skew version, untouched. - * - * The loop structure is dbuf4's (both parents share it), unchanged: - * prologue: prefetch chunk 0 -> temp, store to slice 0 (no barrier) - * per iter: barrier -> prefetch(next) -> MMA(cur) -> store(next) - * kept nested (groups x chunks) with an unconditional group epilog -- - * flattening it crashes the Xclipse PAL compiler at large spec-resolved trip - * counts (see dbuf2's own header). - * - * Selected via - * ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4zpgtr_txkgs<32|64> - * (QuantizedLinear.cpp), additive to the tsweep_dbuf4zpg_t..., tsweep_dbuf4tr_t... - * and tsweep_t... namespaces. NOT the default -- unvalidated until it passes - * repeated test_llama_microbench --correctness-only runs (see - * dq8ca_coopmat_variant()'s comment on why a single pass is not proof). - * - * Performs: out[M,N] = dequant(int8_act) * dequant(int_w) (+ bias) - * via coopmat x coopmat -> coopmat on the matrix unit. - * - * Hard preconditions (dbuf4zpg's, plus dbuf4tr's row-major/alignment ones): - * M % WG_TILE_M == 0, N % WG_TILE_N == 0, K % WG_TILE_K == 0, - * group_size % WG_TILE_K == 0, K % 4 == 0, - * WG_TILE_M % MMA_M == 0, WG_TILE_K % MMA_K == 0, - * t_packed_int8_input in kPackedInt8_4W (row-major) layout, - * device exposes coopmatx-> at 16x16x16. - */ - -#version 450 core - -#extension GL_KHR_cooperative_matrix : require -#extension GL_KHR_memory_scope_semantics : require -#extension GL_KHR_shader_subgroup_basic : enable -#extension GL_EXT_shader_explicit_arithmetic_types : require -#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require -// 8-bit SSBO access: A is bound as a scalar int8_t array so that the -// coopMatLoad below has a MATCHING component type (see dbuf4tr's header for -// why the type must match on this driver). -#extension GL_EXT_shader_8bit_storage : require -#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require -#extension GL_EXT_control_flow_attributes : enable - -#define PRECISION ${PRECISION} - -$if WEIGHT_NBITS == 4: - #define WEIGHT_INT4 - -$if HAS_BIAS: - #define HAS_BIAS - -$if WEIGHT_STORAGE == "buffer": - #define WEIGHT_BUFFER - -$if IO_STORAGE == "texture3d": - #define IO_TEXTURE - -layout(std430) buffer; - -#include "common.glslh" - -// Bindings — match add_linear_dqa_qw_node arg order: -// output(0), fp_input(1), packed_int8_input(2), int_input_sums(3 - unused), -// input_scales(4), input_zps(5), packed_weight(6), weight_sums(7), -// weight_scales(8), bias(9). -${layout_declare_tensor(B, "w", "t_output", "half", IO_STORAGE, is_scalar_array=True)} -// t_input is unread here -- the activations arrive already quantized in -// t_packed_int8_input -- but stays declared so the binding layout matches the -// dispatch site. It tracks IO_STORAGE so the two IO tensors stay consistent. -${layout_declare_tensor(B, "r", "t_input", "half", IO_STORAGE, is_scalar_array=False)} -// ROW-MAJOR (kPackedInt8_4W) packed activations, bound as a scalar int8_t -// array (row stride = K int8) -- dbuf4tr's binding, unchanged. The stock -// 4h4w layout dbuf4zpg uses is NOT row-major (component index selects a row, -// non-affine), so it cannot be addressed by any coopMatLoad. -${layout_declare_tensor(B, "r", "t_packed_int8_input", "int8", "buffer", is_scalar_array=True)} -${layout_declare_tensor(B, "r", "t_int8_input_sums", "int", "buffer", is_scalar_array=True)} -${layout_declare_tensor(B, "r", "t_int8_input_scales", "half", "texture3d")} -${layout_declare_tensor(B, "r", "t_int8_input_zps", "int8", "texture3d")} -${layout_declare_tensor(B, "r", "t_packed_weight", "int", WEIGHT_STORAGE, is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_weight_sums", "int", "buffer", is_scalar_array=True)} -${layout_declare_tensor(B, "r", "t_weight_scales", "half", "buffer", is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_bias", "half", "buffer", is_scalar_array=True)} - -${layout_declare_ubo(B, "ivec4", "output_sizes")} -${layout_declare_ubo(B, "ivec4", "input_sizes")} - -layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; - -${layout_declare_spec_const(C, "int", "apply_bias", "0")} -// INT4 only; inert (0) for INT8 so the dispatcher's spec list lines up. -${layout_declare_spec_const(C, "int", "K4_per_group", "0")} -${layout_declare_spec_const(C, "int", "num_groups_arg", "0")} -${layout_declare_spec_const(C, "int", "out_N_arg", "0")} - -// Tile geometry -const uint MMA_M = ${MMA_M}; -const uint MMA_N = ${MMA_N}; -const uint MMA_K = ${MMA_K}; - -const uint WG_TILE_M = ${WG_TILE_M}; -const uint WG_TILE_N = ${WG_TILE_N}; -const uint WG_TILE_K = ${WG_TILE_K}; - -const uint SG_GRID_X = ${SG_GRID_X}; -const uint SG_GRID_Y = ${SG_GRID_Y}; -const uint SUBGROUP_SIZE = ${SUBGROUP_SIZE}; -const uint NUM_SUBGROUPS = SG_GRID_X * SG_GRID_Y; -const uint WG_SIZE = NUM_SUBGROUPS * SUBGROUP_SIZE; - -const uint SG_TILE_M = WG_TILE_M / SG_GRID_Y; -const uint SG_TILE_N = WG_TILE_N / SG_GRID_X; -const uint MMAS_PER_SG_M = SG_TILE_M / MMA_M; -const uint MMAS_PER_SG_N = SG_TILE_N / MMA_N; - -const uint A_SLAB_INT8 = WG_TILE_M * MMA_K; -const uint B_USEFUL_U32 = MMA_K / 4u; -// No skew + coalesced write -- dbuf4zpg's B fix, unchanged (this file does -// not touch B staging at all). -const uint B_STRIDE_U32 = B_USEFUL_U32; -const uint B_SLAB_U32 = WG_TILE_N * B_STRIDE_U32; -const uint NUM_K_SLABS = WG_TILE_K / MMA_K; - -const uint A_SLAB_U32 = A_SLAB_INT8 / 4u; -const uint A_STRIDE_U32 = MMA_K / 4u; - -// One ping-pong slice covers all K-slabs of one chunk. -const uint ASH_SLICE_U32 = NUM_K_SLABS * A_SLAB_U32; -const uint BSH_SLICE_U32 = NUM_K_SLABS * B_SLAB_U32; - -// Double-buffered MMA operand staging. -shared uint Ash_int8[2u * ASH_SLICE_U32]; -// ===== bw3: shared B retyped to uvec2 ===== -// The ONLY combination that gets both a genuinely wide LDS store and the -// safe ownership. A uvec4 element would be a whole column-slab, forcing -// 4-wide-per-thread ownership with half the workgroup idle -- which is exactly -// the pattern that fails the num_groups==2 shapes (see dbuf4zpgbw/bwr/bv4). -// A uvec2 element is 2 uints, so BSH_SLICE_V2 == WG_SIZE at the shipped -// geometry: one element per thread, ALL 256 threads active (dbuf4zpgbw2's -// proven-correct ownership), one ds_write_b64 each. -// Packing ratio for the compute-side coopMatLoad is 8x (int8 coopmat from -// uvec2[]); dbuf4zpgbw2 showed adjacent scalar writes do NOT get merged by the -// compiler, so retyping is the only way to actually emit a wide store. -const uint B_STRIDE_V2 = B_STRIDE_U32 / 2u; // uvec2 per column -const uint B_SLAB_V2 = B_SLAB_U32 / 2u; -const uint BSH_SLICE_V2 = BSH_SLICE_U32 / 2u; -const uint V2_PER_THREAD = BSH_SLICE_V2 / WG_SIZE; -shared uvec2 Bsh_v2[2u * BSH_SLICE_V2]; - -// Per-WG-tile-row activation params (loaded ONCE at WG start; constant -// across groups). -shared int izp_sh[WG_TILE_M]; // int32 (cast from int8 source) for broadcast -shared float ifs_sh[WG_TILE_M]; // float32 (cast from fp16 source) for broadcast - -// Per-(group, output-channel) weight params, ping-ponged by group parity. -// (For per-channel INT8 only slice 0 is ever used.) -shared float wsc_sh[2u * WG_TILE_N]; -// SUM_g wsc[g][n]*wsum[g][n] per output channel -- weight-side only, so it is -// accumulated once in the prologue. dbuf4zpg's zp-hoist, unchanged. -shared float wcorr_sh[WG_TILE_N]; - -#ifdef HAS_BIAS -shared float bias_sh[WG_TILE_N]; -#endif - -#ifdef IO_TEXTURE -// Result staging for the imageStore epilogue, mirroring the fp16 kernel: -// SG_GRID_Y bands of MMA_M rows, each WG_TILE_N wide, row-major. A full -// WG_TILE_M x WG_TILE_N buffer would cost SG_GRID_Y/MMAS_PER_SG_M x more LDS -// and wreck occupancy. float16_t-typed because coopMatStore needs it. -const uint CSH_ROWS = SG_GRID_Y * MMA_M; -shared float16_t Csh_out[CSH_ROWS * WG_TILE_N]; -#endif - -// Running fp32 accumulator (across all groups). -coopmat - result[MMAS_PER_SG_M][MMAS_PER_SG_N]; - -// Per-group int32 MMA accumulator. -coopmat - accum_int32[MMAS_PER_SG_M][MMAS_PER_SG_N]; - - -// Byte-parallel int4 -> int8 widening. dbuf4zpg's, unchanged (B-side only). -// -// The four nibbles this shader needs from one packed uint are ALREADY one per -// byte (bits 3:0 of each byte for parity 0, bits 7:4 for parity 1), so all four -// can be widened at once instead of with a per-nibble -// shift/mask/bias-subtract/mask chain. -// -// For v in [0,15] the biased value is v-8. `v ^ 8` is exactly the 4-bit -// two's-complement pattern of v-8, because -8 == +8 (mod 16): -// v=0 -> 0x8 -> -8 v=7 -> 0xF -> -1 -// v=8 -> 0x0 -> 0 v=15 -> 0x7 -> +7 -// so the only remaining work is sign-extending bit 3 into bits 7:4 per byte. -// `sgn * 0x1E` does that with no cross-byte carry: 0x08 * 0x1E == 0xF0 exactly, -// and sgn is at most 0x08080808 so the product is at most 0xF0F0F0F0. -// -// A naive `nib - 0x08080808` would NOT work -- it borrows across byte lanes -// whenever a nibble is < 8. Shifts must be on uint, not int, so the >> is -// logical rather than arithmetic. -// -// ~5 ops per 4 weights vs ~22 for the per-nibble chain; bit-identical output. -uint widen_nibbles(const uint w, const uint parity) { - const uint nib = (parity == 0u) ? (w & 0x0F0F0F0Fu) : ((w >> 4u) & 0x0F0F0F0Fu); - const uint p = nib ^ 0x08080808u; - const uint sgn = p & 0x08080808u; - return p | (sgn * 0x1Eu); -} - -void main() { - const uvec2 tileID = uvec2(gl_WorkGroupID.xy); - const uvec2 warpInTile = uvec2( - gl_SubgroupID % SG_GRID_X, - gl_SubgroupID / SG_GRID_X); - - const uint K = uint(input_sizes.x); - const uint N = uint(output_sizes.x); - const uint N4 = (N + 3u) / 4u; - const uint nblocks_x_A = (K + 3u) >> 2u; - // A row stride in INT8 elements (dbuf4tr's binding is row-major int8, not - // the 4h4w ivec4 block layout dbuf4zpg used -- so A addressing below is in - // int8 elements, not int, and derived from nblocks_x_A so it matches the - // packer's `m_row * K4 + k4` addressing exactly (K % 4 == 0 makes them - // equal to K directly). - const uint a_row_stride_i8 = nblocks_x_A * 4u; - -#ifdef WEIGHT_INT4 - const uint num_groups = uint(num_groups_arg); - const uint CHUNKS_PER_GROUP = uint(K4_per_group) * 4u / WG_TILE_K; -#else - // Per-channel: a single quant "group" spanning all of K. The nested - // groups x chunks loop below collapses to a flat chunk loop, the wsum/wsc - // ping-pong never crosses a boundary, and the epilog runs exactly once. - const uint num_groups = 1u; - const uint CHUNKS_PER_GROUP = uint(num_groups_arg); -#endif - const uint num_chunks = num_groups * CHUNKS_PER_GROUP; - - const uint tile_m_start = WG_TILE_M * tileID.y; - const uint tile_n_start = WG_TILE_N * tileID.x; - - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - result[i][j] = coopmat(0.0); - accum_int32[i][j] = coopmat(0); - } - } - - const uint K_BLOCKS_PER_CHUNK = WG_TILE_K >> 2u; - - // --- A staging tile map (dbuf4tr's, unmodified): one MMA_M x MMA_K coopmat - // tile per subgroup per slot, dealt round-robin across the - // NUM_SUBGROUPS subgroups so every subgroup participates. Replaces - // dbuf4zpg's per-thread (m4, k4) map / a_active guard entirely -- see - // design.md D3 for why this is reused as-is, not re-derived. - const uint A_TILES_M = WG_TILE_M / MMA_M; - const uint A_TILES_K = WG_TILE_K / MMA_K; // == NUM_K_SLABS - const uint NUM_A_TILES = A_TILES_M * A_TILES_K; - const uint A_TILES_PER_SG = (NUM_A_TILES + NUM_SUBGROUPS - 1u) / NUM_SUBGROUPS; - -#ifdef WEIGHT_INT4 - // --- B staging thread map: (block, col) slots; each slot extracts one - // ColumnMajor LDS uint (4 K-contiguous sign-extended int8) --- - const uint B_TOTAL_SLOTS = K_BLOCKS_PER_CHUNK * WG_TILE_N; - const uint B_SLOTS_PER_THREAD = B_TOTAL_SLOTS / WG_SIZE; - const uint N8_PER_TILE = WG_TILE_N >> 3u; -#else - // --- B staging thread map: one (k4, n4) ivec4 block per active thread --- - const uint B_FETCH_SLOTS = K_BLOCKS_PER_CHUNK * (WG_TILE_N >> 2u); - const uint N4_PER_TILE = WG_TILE_N >> 2u; - const uint b_k4_in_chunk = gl_LocalInvocationID.x / N4_PER_TILE; - const uint b_n_uint_col = gl_LocalInvocationID.x % N4_PER_TILE; - const bool b_active = gl_LocalInvocationID.x < B_FETCH_SLOTS; -#endif - - // ===== INTERVENTION F: hoist loop-invariant B staging index math ===== - // dbuf4zpg's, unchanged -- B staging is untouched by this file's A-staging - // swap. See dbuf4zpg's header for the full rationale (ablation-attributed - // -16.8% block, mostly index arithmetic not memory traffic). -#ifdef WEIGHT_INT4 - // bw3 ownership == dbuf4zpgbw2's (proven correct), one uvec2 element each. - uint b_v2off[V2_PER_THREAD]; // uvec2 element index within a slice - uint b_comp1[V2_PER_THREAD]; - uint b_par1[V2_PER_THREAD]; - uint b_n8b[V2_PER_THREAD]; - uint b_k4b[V2_PER_THREAD]; // k4 of the FIRST uint of the pair - [[unroll]] for (uint si = 0; si < V2_PER_THREAD; ++si) { - const uint e = gl_LocalInvocationID.x + si * WG_SIZE; // uvec2 index - const uint u0 = e * 2u; // first uint - const uint slab_idx = u0 / B_SLAB_U32; - const uint inslab = u0 % B_SLAB_U32; - const uint n_col = inslab / B_STRIDE_U32; - const uint k4_lo = inslab % B_STRIDE_U32; - const uint rem = n_col & 7u; - b_v2off[si] = e; - b_comp1[si] = rem & 3u; - b_par1[si] = rem >> 2u; - b_n8b[si] = (tile_n_start >> 3u) + (n_col >> 3u); - b_k4b[si] = slab_idx * (MMA_K >> 2u) + k4_lo; - } -#endif - - // Prefetch temp registers. temp_A is a coopmat array (dbuf4tr's A-staging - // technique); indices into it are [[unroll]]-resolved compile-time - // constants, never dynamic -- dynamic indexing of a coopmat array is - // exactly the construct the Xclipse/AMD-PAL compiler has miscompiled - // before. - coopmat - temp_A[A_TILES_PER_SG]; -#ifdef WEIGHT_INT4 - ivec4 temp_B[V2_PER_THREAD * 2u]; // bw3: 2 packed fetches per uvec2 - float temp_wsc; -#else - ivec4 temp_B; -#endif - - // ========================================================= - // PROLOGUE - // ========================================================= - if (gl_LocalInvocationID.x < (WG_TILE_M >> 2u)) { - const uint m4 = (tile_m_start >> 2u) + gl_LocalInvocationID.x; - const vec4 sc = vec4(texelFetch(t_int8_input_scales, ivec3(m4, 0, 0), 0)); - const ivec4 zp = texelFetch(t_int8_input_zps, ivec3(m4, 0, 0), 0); - const uint base = gl_LocalInvocationID.x * 4u; - ifs_sh[base + 0u] = sc.x; ifs_sh[base + 1u] = sc.y; - ifs_sh[base + 2u] = sc.z; ifs_sh[base + 3u] = sc.w; - izp_sh[base + 0u] = zp.x; izp_sh[base + 1u] = zp.y; - izp_sh[base + 2u] = zp.z; izp_sh[base + 3u] = zp.w; - } - // Group 0 weight scales -> slice 0, and the hoisted weight-side correction - // SUM_g wsc[g][n]*wsum[g][n] accumulated across ALL groups. dbuf4zpg's - // zp-hoist, unchanged. - if (gl_LocalInvocationID.x < WG_TILE_N) { - const uint n_idx = tile_n_start + gl_LocalInvocationID.x; - f16vec4 sv0 = t_weight_scales[n_idx >> 2u]; - wsc_sh[gl_LocalInvocationID.x] = float(sv0[n_idx & 3u]); - - float corr = 0.0; - for (uint g = 0; g < num_groups; ++g) { - f16vec4 sv = t_weight_scales[g * N4 + (n_idx >> 2u)]; - corr += float(sv[n_idx & 3u]) * float(t_weight_sums[g * N + n_idx]); - } - wcorr_sh[gl_LocalInvocationID.x] = corr; - } - memoryBarrierShared(); - barrier(); - - // NOTE: dbuf4zpg builds izp_bcast/ifs_bcast AFTER the group loop, not here - // -- that is the register-pressure saving zp-hoist buys. Unchanged. - - // dbuf4: prefetch chunk 0 into temp registers, THEN store to slice 0 (no - // barrier here -- the main loop's first iteration barriers before - // reading slice 0). - // - // A staging (dbuf4tr's technique): per-subgroup coopMatLoad straight from - // the row-major global buffer. - [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { - const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; - if (t < NUM_A_TILES) { - const uint tm = t / A_TILES_K; - const uint tk = t % A_TILES_K; - coopMatLoad( - temp_A[s], t_packed_int8_input, - (tile_m_start + tm * MMA_M) * a_row_stride_i8 + tk * MMA_K, - a_row_stride_i8, - gl_CooperativeMatrixLayoutRowMajor); - } - } -#ifdef WEIGHT_INT4 - [[unroll]] for (uint si = 0; si < V2_PER_THREAD; ++si) { - [[unroll]] for (uint h = 0; h < 2u; ++h) { -#ifdef WEIGHT_BUFFER - temp_B[si * 2u + h] = t_packed_weight[(b_n8b[si] * nblocks_x_A) + b_k4b[si] + h]; -#else - temp_B[si * 2u + h] = texelFetch(t_packed_weight, ivec2(b_k4b[si] + h, b_n8b[si]), 0); -#endif - } - } -#else - if (b_active) { - const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; -#ifdef WEIGHT_BUFFER - temp_B = t_packed_weight[(b_k4_in_chunk * N4) + block_x_w]; -#else - temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, b_k4_in_chunk), 0); -#endif - } -#endif - { - // store chunk 0 -> slice 0 - // A staging (dbuf4tr's technique): coopMatStore into the same Ash_int8 - // slot layout dbuf4zpg's scalar scatter used to write. - [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { - const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; - if (t < NUM_A_TILES) { - const uint tm = t / A_TILES_K; - const uint tk = t % A_TILES_K; - coopMatStore( - temp_A[s], Ash_int8, - tk * A_SLAB_U32 + (tm * MMA_M) * A_STRIDE_U32, - A_STRIDE_U32, - gl_CooperativeMatrixLayoutRowMajor); - } - } -#ifdef WEIGHT_INT4 - [[unroll]] for (uint si = 0; si < V2_PER_THREAD; ++si) { - // ONE ds_write_b64 per uvec2 element. - Bsh_v2[ b_v2off[si]] = uvec2( - widen_nibbles(uint(temp_B[si * 2u + 0u][b_comp1[si]]), b_par1[si]), - widen_nibbles(uint(temp_B[si * 2u + 1u][b_comp1[si]]), b_par1[si])); - } -#else - if (b_active) { - const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); - const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); - const uint n_col_base = b_n_uint_col * 4u; - [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { - // bw3: same address as (uvec2 element, component). This INT8 path - // writes 4 different columns for one k4, so no wide store applies. - { - const uint u = slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab; - Bsh_v2[u / 2u][u % 2u] = uint(temp_B[n_in_blk]); - } - } - } -#endif - } - - // ========================================================= - // MAIN LOOP (dbuf4) — nested groups x chunks (kept nested; flattening it - // with a conditional coopmat epilog crashes the Xclipse PAL compiler at - // large spec-resolved trip counts). One barrier per chunk. Chunk - // iteration (global index `chunk`): - // 1. barrier — A/B slice (chunk%2) fully written; on the first chunk - // of group g, wsc slice (g%2) is too. - // 2. prefetch — chunk+1 (A tiles, B blocks) into temp; when chunk+1 - // starts a new group, also its wsc element. Skipped - // entirely on the final chunk. - // 3. int8 MMA — on slice (chunk%2) into accum_int32. - // 4. store — temp -> A/B slice ((chunk+1)%2), unpacking the weight; - // on a group boundary, wsc -> slice ((g+1)%2). - // The group epilog runs unconditionally at the tail of each group. - // ========================================================= - uint chunk = 0; - for (uint group_i = 0; group_i < num_groups; ++group_i) { - for (uint inner = 0; inner < CHUNKS_PER_GROUP; ++inner, ++chunk) { - const bool has_next = chunk + 1u < num_chunks; - const bool group_crossing = has_next && (inner + 1u == CHUNKS_PER_GROUP); - const uint cur_a = (chunk % 2u) * ASH_SLICE_U32; - const uint cur_b_v2 = (chunk % 2u) * BSH_SLICE_V2; - const uint nxt_a = ((chunk + 1u) % 2u) * ASH_SLICE_U32; - const uint nxt_b_v2 = ((chunk + 1u) % 2u) * BSH_SLICE_V2; - - // coopmat-lds-fence: barrier() alone does NOT order shared stores against a - // subsequent coopMatLoad on the M51 Xclipse/AMD-PAL driver -- symptom is one - // stale MMA_M-row band of A, all columns, ~2.5% of runs, no crash (observed - // 2026-09-02 in sdpa_compute_out_coopmat.glsl). Measured cost of the fence: - // none (see this change's results). See memory - // `coopmat-lds-needs-explicit-memorybarriershared`. - memoryBarrierShared(); - barrier(); - - // --- 2. prefetch chunk+1 -> temp --- - if (has_next) { - const uint chunkK_nxt = (chunk + 1u) * WG_TILE_K; - // A staging (dbuf4tr's technique): coopMatLoad straight from global. - [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { - const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; - if (t < NUM_A_TILES) { - const uint tm = t / A_TILES_K; - const uint tk = t % A_TILES_K; - coopMatLoad( - temp_A[s], t_packed_int8_input, - (tile_m_start + tm * MMA_M) * a_row_stride_i8 + chunkK_nxt + - tk * MMA_K, - a_row_stride_i8, - gl_CooperativeMatrixLayoutRowMajor); - } - } -#ifdef WEIGHT_INT4 - [[unroll]] for (uint si = 0; si < V2_PER_THREAD; ++si) { - [[unroll]] for (uint h = 0; h < 2u; ++h) { - const uint k4_blk = (chunkK_nxt >> 2u) + b_k4b[si] + h; -#ifdef WEIGHT_BUFFER - temp_B[si * 2u + h] = t_packed_weight[(b_n8b[si] * nblocks_x_A) + k4_blk]; -#else - temp_B[si * 2u + h] = texelFetch(t_packed_weight, ivec2(k4_blk, b_n8b[si]), 0); -#endif - } - } - if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { - const uint n_idx = tile_n_start + gl_LocalInvocationID.x; - f16vec4 sv = t_weight_scales[(group_i + 1u) * N4 + (n_idx >> 2u)]; - temp_wsc = float(sv[n_idx & 3u]); - } -#else - if (b_active) { - const uint block_y_w = (chunkK_nxt >> 2u) + b_k4_in_chunk; - const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; -#ifdef WEIGHT_BUFFER - temp_B = t_packed_weight[(block_y_w * N4) + block_x_w]; -#else - temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, block_y_w), 0); -#endif - } -#endif - } - - // --- 3. int8 MMA on the cur slice --- - [[unroll]] for (uint k = 0; k < NUM_K_SLABS; ++k) { - const uint slab_a_base_u32 = cur_a + k * A_SLAB_U32; - const uint slab_b_base_v2 = cur_b_v2 + k * B_SLAB_V2; - - coopmat matA[MMAS_PER_SG_M]; - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - const uint row_a = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); - coopMatLoad( - matA[i], Ash_int8, - slab_a_base_u32 + row_a * A_STRIDE_U32, - A_STRIDE_U32, - gl_CooperativeMatrixLayoutRowMajor); - } - - coopmat matB; - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint col_b = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopMatLoad( - matB, Bsh_v2, - slab_b_base_v2 + col_b * B_STRIDE_V2, - B_STRIDE_V2, - gl_CooperativeMatrixLayoutColumnMajor); - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - accum_int32[i][j] = coopMatMulAdd(matA[i], matB, accum_int32[i][j]); - } - } - } - - // --- 4. store temp (chunk+1) -> nxt slice --- - if (has_next) { - // A staging (dbuf4tr's technique): coopMatStore into the nxt slice. - [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { - const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; - if (t < NUM_A_TILES) { - const uint tm = t / A_TILES_K; - const uint tk = t % A_TILES_K; - coopMatStore( - temp_A[s], Ash_int8, - nxt_a + tk * A_SLAB_U32 + (tm * MMA_M) * A_STRIDE_U32, - A_STRIDE_U32, - gl_CooperativeMatrixLayoutRowMajor); - } - } -#ifdef WEIGHT_INT4 - [[unroll]] for (uint si = 0; si < V2_PER_THREAD; ++si) { - // ONE ds_write_b64 per uvec2 element. - Bsh_v2[nxt_b_v2 + b_v2off[si]] = uvec2( - widen_nibbles(uint(temp_B[si * 2u + 0u][b_comp1[si]]), b_par1[si]), - widen_nibbles(uint(temp_B[si * 2u + 1u][b_comp1[si]]), b_par1[si])); - } - if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { - const uint wbase_nxt = ((group_i + 1u) % 2u) * WG_TILE_N; - wsc_sh[wbase_nxt + gl_LocalInvocationID.x] = temp_wsc; - } -#else - if (b_active) { - const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); - const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); - const uint n_col_base = b_n_uint_col * 4u; - [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { - { - const uint u = slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab; - Bsh_v2[nxt_b_v2 + u / 2u][u % 2u] = uint(temp_B[n_in_blk]); - } - } - } -#endif - } - } // chunks - - // --- Group epilog: scale-only accumulate, reset accum --- - // dbuf4zpg's, unchanged. Just result += float(acc) * wsc. The - // zero-point subtract and the ifs multiply are hoisted out of the group - // loop (applied once below). - { - const uint wbase = (group_i % 2u) * WG_TILE_N; - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - - coopmat wsc_bcast; - coopMatLoad( - wsc_bcast, wsc_sh, - wbase + local_n_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutRowMajor); - - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - result[i][j] += - coopmat( - accum_int32[i][j]) * wsc_bcast; - accum_int32[i][j] = coopmat(0); - } - } - } - } // groups - - // --- Hoisted correction, applied ONCE: -------------------------------- - // result = ifs * ( result - izp * SUM_g wsc_g*wsum_g ) - // dbuf4zpg's, unchanged. izp/ifs are loaded here rather than before the - // group loop so they are not live across it. - { - coopmat - izpf_bcast[MMAS_PER_SG_M]; - coopmat - ifs_bcast[MMAS_PER_SG_M]; - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - const uint local_m_base = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); - coopmat izp_i; - coopMatLoad( - izp_i, izp_sh, local_m_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutColumnMajor); - izpf_bcast[i] = - coopmat(izp_i); - coopMatLoad( - ifs_bcast[i], ifs_sh, local_m_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutColumnMajor); - } - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopmat wcorr_bcast; - coopMatLoad( - wcorr_bcast, wcorr_sh, local_n_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutRowMajor); - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - result[i][j] = ifs_bcast[i] * (result[i][j] - izpf_bcast[i] * wcorr_bcast); - } - } - } - - // --- Bias (optional) --- -#ifdef HAS_BIAS - if (apply_bias > 0) { - for (uint t = gl_LocalInvocationID.x; t < WG_TILE_N; t += WG_SIZE) { - bias_sh[t] = float(t_bias[tile_n_start + t]); - } - memoryBarrierShared(); - barrier(); - } -#endif - - // --- Store result tile --- - // N for the store address math MUST come from the spec constant, not the - // sizes UBO (see out_N_arg above). -#ifdef IO_TEXTURE - // Epilogue iteration i drains accumulator row-block i from EVERY subgroup - // into Csh_out at once, so the SG_GRID_Y bands it holds are disjoint global - // row ranges; the whole workgroup then imageStores them. lr / MMA_M is the - // writing subgroup's warpInTile.y, so the global row reproduces the buffer - // path's gi exactly. - // - // PORTABILITY NOTE: the barrier() in the loop body keeps this loop rolled - // despite [[unroll]], so result[i][j] IS dynamically indexed. Coopmat arrays - // are opaque per-lane storage and dynamic indexing is exactly the construct - // the Xclipse/AMD-PAL compiler has broken before -- check this first if the - // texture variants miscompile on M51. - const uint CSH_TEXELS_PER_ROW = WG_TILE_N / 4u; - const uint CSH_TEXELS = CSH_ROWS * CSH_TEXELS_PER_ROW; - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - // Guards Csh_out against the previous iteration's readers. Inert on i == 0 - // but must stay unconditional to remain workgroup-uniform. - // coopmat-lds-fence: barrier() alone does NOT order shared stores against a - // subsequent coopMatLoad on the M51 Xclipse/AMD-PAL driver -- symptom is one - // stale MMA_M-row band of A, all columns, ~2.5% of runs, no crash (observed - // 2026-09-02 in sdpa_compute_out_coopmat.glsl). Measured cost of the fence: - // none (see this change's results). See memory - // `coopmat-lds-needs-explicit-memorybarriershared`. - memoryBarrierShared(); - barrier(); - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { -#ifdef HAS_BIAS - if (apply_bias > 0) { - const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopmat bias_tile; - coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); - result[i][j] += bias_tile; - } -#endif - coopmat out_tile = - coopmat(result[i][j]); - coopMatStore( - out_tile, Csh_out, - warpInTile.y * MMA_M * WG_TILE_N + - MMA_N * (MMAS_PER_SG_N * warpInTile.x + j), - WG_TILE_N, - gl_CooperativeMatrixLayoutRowMajor); - } - memoryBarrierShared(); - barrier(); - - for (uint t = gl_LocalInvocationID.x; t < CSH_TEXELS; t += WG_SIZE) { - const uint lr = t / CSH_TEXELS_PER_ROW; - const uint lc4 = t % CSH_TEXELS_PER_ROW; - const uint m = - tile_m_start + (lr / MMA_M) * SG_TILE_M + i * MMA_M + (lr % MMA_M); - const uint base = lr * WG_TILE_N + lc4 * 4u; - imageStore( - t_output, - ivec3(tile_n_start / 4u + lc4, m, 0), - vec4( - float(Csh_out[base]), - float(Csh_out[base + 1u]), - float(Csh_out[base + 2u]), - float(Csh_out[base + 3u]))); - } - } -#else - const uint N_out = uint(out_N_arg); - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint gi = tile_m_start + MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); - const uint gj = tile_n_start + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - -#ifdef HAS_BIAS - if (apply_bias > 0) { - const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopmat bias_tile; - coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); - result[i][j] += bias_tile; - } -#endif - - coopmat out_tile = - coopmat(result[i][j]); - coopMatStore( - out_tile, t_output, - gi * N_out + gj, N_out, - gl_CooperativeMatrixLayoutRowMajor); - } - } -#endif // IO_TEXTURE -} diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw3.yaml b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw3.yaml deleted file mode 100644 index d167ca95809..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw3.yaml +++ /dev/null @@ -1,60 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -# "zpg" + "-tr" combination: tsweep_dbuf4zpg with its per-thread scalar -# A-staging replaced by tsweep_dbuf4tr's coopMat-mediated A staging. B -# staging, zp-hoist, byte-parallel nibble widening are all dbuf4zpg's, -# unchanged -- only A-staging differs. Requires t_packed_int8_input in the -# ROW-MAJOR kPackedInt8_4W layout (same requirement as tsweep_dbuf4tr). -# -# PROMOTED 2026-09-01 as the shipped default -- see QuantizedLinear.cpp's -# dq8ca_coopmat_variant() for the full validation record (10/10 buffer, 6/6 -# texture3d correctness across 1B/3B/8B; +4.2% real e2e prefill; 46.50% -> -# 49.94% efficiency on 8B). Also selectable explicitly via -# ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4zpgbw3_txkgs. -# See openspec/changes/archive/2026-08-31-dq8ca-tr-staged-a-on-zpg. -# -# Single seed tile: the prior dbuf4zpg default's own tile (t128x64k32g42s32). -# A re-sweep against this shader's own (lower) register-pressure profile was -# run as a follow-up (coopmat-tr-tilesweep-4w-port) and found no better tile -# -- this remains the best known geometry. - -linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw3: - parameter_names_with_default_values: - PRECISION: highp - HAS_BIAS: false - IO_STORAGE: buffer - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - MMA_M: 16 - MMA_N: 16 - MMA_K: 16 - WG_TILE_M: 64 - WG_TILE_N: 32 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - shader_variants: - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw3_t128x64k32g42s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbw3_t128x64k32g42s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbwr.glsl b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbwr.glsl deleted file mode 100644 index e187ae60d97..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbwr.glsl +++ /dev/null @@ -1,783 +0,0 @@ -/* - * Copyright (c) Meta Platforms, Inc. and affiliates. - * All rights reserved. - * - * This source code is licensed under the BSD-style license found in the - * LICENSE file in the root directory of this source tree. - */ - -/* - * "zpg" + "-tr" combination: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpg.glsl - * with its per-thread scalar A-staging replaced by - * linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr.glsl's coopMat-mediated A - * staging, PLUS "bw": B staging ownership moved from (block,col) slots spread - * over 4 threads to ONE thread per (K-slab, column), writing its 4 uints at - * consecutive LDS indices so they are merge-eligible. Shared-array types, the - * compute-side loads, the quantization math and every safeguard are - * dbuf4zpgtr's, unchanged. This is - * an ADDITIVE combination, not a redesign: every non-A-staging block below - * (B staging: coalesced write, no skew; zp-hoist: izp/ifs applied once after - * the group loop via wcorr_sh; byte-parallel nibble widening; static - * A_ALWAYS_ACTIVE branch elision -- N/A here, see below; group epilog; - * bias/store epilogue) is byte-identical to dbuf4zpg's. Only the A-staging - * block (prologue load+store, main-loop prefetch+store) is dbuf4tr's, - * verbatim. - * - * dbuf4zpg's per-thread A staging used an `a_active` guard (statically always - * true when A_ACTIVE_THREADS == WG_SIZE, via the A_MAP_FULL-gated - * A_ALWAYS_ACTIVE macro). dbuf4tr's per-SUBGROUP tile map has no equivalent - * concept -- every subgroup participates via a `t < NUM_A_TILES` guard that - * depends only on gl_SubgroupID, not gl_LocalInvocationID.x -- so - * A_MAP_FULL/A_ALWAYS_ACTIVE is dropped entirely in this file; it would be - * dead code for the new A-staging block. - * - * Rationale for combining this way (not the reverse) and why this is worth - * building at all: see this change's design.md D0-D3. In short -- the only - * existing measurement of dbuf4tr's A-staging technique (28.72-30.51%, - * dq8ca-arch-redesign) was taken against dbuf4tr's own pre-zpg baseline - * (old B skew, no byte-parallel widening, no branch elision) -- a materially - * weaker shader than the 46.49-46.50% dbuf4zpg this file now combines it - * with. This file exists to answer whether that combination performs - * differently now that register pressure is already reduced. - * - * A staging (the actual delta from dbuf4zpg): - * dbuf4zpg: per-thread (m4, k4) ivec4 fetch, hoisted a_lds_off0/a_glb_row; - * only A_ACTIVE_THREADS invocations participate, each scattering - * 4 rows into Ash_int8 with 4 scalar stores. - * this file: per-SUBGROUP MMA_M x MMA_K tile fetch via coopMatLoad straight - * from a ROW-MAJOR (kPackedInt8_4W) int8 activation buffer, then - * coopMatStore into the same Ash_int8 slot -- dbuf4tr's mapping, - * unmodified (not re-derived; see design.md D3). - * - * t_packed_int8_input is therefore bound the same way dbuf4tr binds it: a - * SCALAR int8_t array in the kPackedInt8_4W layout (plain row-major int8, - * row stride K), produced by quantize_and_pack_4w_with_group_sums.glsl. - * QuantizedLinear.cpp's dq8ca_variant_wants_rowmajor_a() must recognize this - * file's variant token (tsweep_dbuf4zpgtr_t...) the same way it already - * recognizes tsweep_dbuf4tr_t/trm_t/trd_t, so graph-build time (packer - * selection) and dispatch time (kernel selection) cannot disagree. - * - * B CANNOT be coopmat-staged (int4 nibble unpack; a coopmat's per-lane layout - * is opaque to hand-assembly from unpacked registers) -- unchanged from both - * parent files. B staging below is dbuf4zpg's byte-parallel, coalesced, - * no-skew version, untouched. - * - * The loop structure is dbuf4's (both parents share it), unchanged: - * prologue: prefetch chunk 0 -> temp, store to slice 0 (no barrier) - * per iter: barrier -> prefetch(next) -> MMA(cur) -> store(next) - * kept nested (groups x chunks) with an unconditional group epilog -- - * flattening it crashes the Xclipse PAL compiler at large spec-resolved trip - * counts (see dbuf2's own header). - * - * Selected via - * ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4zpgtr_txkgs<32|64> - * (QuantizedLinear.cpp), additive to the tsweep_dbuf4zpg_t..., tsweep_dbuf4tr_t... - * and tsweep_t... namespaces. NOT the default -- unvalidated until it passes - * repeated test_llama_microbench --correctness-only runs (see - * dq8ca_coopmat_variant()'s comment on why a single pass is not proof). - * - * Performs: out[M,N] = dequant(int8_act) * dequant(int_w) (+ bias) - * via coopmat x coopmat -> coopmat on the matrix unit. - * - * Hard preconditions (dbuf4zpg's, plus dbuf4tr's row-major/alignment ones): - * M % WG_TILE_M == 0, N % WG_TILE_N == 0, K % WG_TILE_K == 0, - * group_size % WG_TILE_K == 0, K % 4 == 0, - * WG_TILE_M % MMA_M == 0, WG_TILE_K % MMA_K == 0, - * t_packed_int8_input in kPackedInt8_4W (row-major) layout, - * device exposes coopmatx-> at 16x16x16. - */ - -#version 450 core - -#extension GL_KHR_cooperative_matrix : require -#extension GL_KHR_memory_scope_semantics : require -#extension GL_KHR_shader_subgroup_basic : enable -#extension GL_EXT_shader_explicit_arithmetic_types : require -#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require -// 8-bit SSBO access: A is bound as a scalar int8_t array so that the -// coopMatLoad below has a MATCHING component type (see dbuf4tr's header for -// why the type must match on this driver). -#extension GL_EXT_shader_8bit_storage : require -#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require -#extension GL_EXT_control_flow_attributes : enable - -#define PRECISION ${PRECISION} - -$if WEIGHT_NBITS == 4: - #define WEIGHT_INT4 - -$if HAS_BIAS: - #define HAS_BIAS - -$if WEIGHT_STORAGE == "buffer": - #define WEIGHT_BUFFER - -$if IO_STORAGE == "texture3d": - #define IO_TEXTURE - -layout(std430) buffer; - -#include "common.glslh" - -// Bindings — match add_linear_dqa_qw_node arg order: -// output(0), fp_input(1), packed_int8_input(2), int_input_sums(3 - unused), -// input_scales(4), input_zps(5), packed_weight(6), weight_sums(7), -// weight_scales(8), bias(9). -${layout_declare_tensor(B, "w", "t_output", "half", IO_STORAGE, is_scalar_array=True)} -// t_input is unread here -- the activations arrive already quantized in -// t_packed_int8_input -- but stays declared so the binding layout matches the -// dispatch site. It tracks IO_STORAGE so the two IO tensors stay consistent. -${layout_declare_tensor(B, "r", "t_input", "half", IO_STORAGE, is_scalar_array=False)} -// ROW-MAJOR (kPackedInt8_4W) packed activations, bound as a scalar int8_t -// array (row stride = K int8) -- dbuf4tr's binding, unchanged. The stock -// 4h4w layout dbuf4zpg uses is NOT row-major (component index selects a row, -// non-affine), so it cannot be addressed by any coopMatLoad. -${layout_declare_tensor(B, "r", "t_packed_int8_input", "int8", "buffer", is_scalar_array=True)} -${layout_declare_tensor(B, "r", "t_int8_input_sums", "int", "buffer", is_scalar_array=True)} -${layout_declare_tensor(B, "r", "t_int8_input_scales", "half", "texture3d")} -${layout_declare_tensor(B, "r", "t_int8_input_zps", "int8", "texture3d")} -${layout_declare_tensor(B, "r", "t_packed_weight", "int", WEIGHT_STORAGE, is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_weight_sums", "int", "buffer", is_scalar_array=True)} -${layout_declare_tensor(B, "r", "t_weight_scales", "half", "buffer", is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_bias", "half", "buffer", is_scalar_array=True)} - -${layout_declare_ubo(B, "ivec4", "output_sizes")} -${layout_declare_ubo(B, "ivec4", "input_sizes")} - -layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; - -${layout_declare_spec_const(C, "int", "apply_bias", "0")} -// INT4 only; inert (0) for INT8 so the dispatcher's spec list lines up. -${layout_declare_spec_const(C, "int", "K4_per_group", "0")} -${layout_declare_spec_const(C, "int", "num_groups_arg", "0")} -${layout_declare_spec_const(C, "int", "out_N_arg", "0")} - -// Tile geometry -const uint MMA_M = ${MMA_M}; -const uint MMA_N = ${MMA_N}; -const uint MMA_K = ${MMA_K}; - -const uint WG_TILE_M = ${WG_TILE_M}; -const uint WG_TILE_N = ${WG_TILE_N}; -const uint WG_TILE_K = ${WG_TILE_K}; - -const uint SG_GRID_X = ${SG_GRID_X}; -const uint SG_GRID_Y = ${SG_GRID_Y}; -const uint SUBGROUP_SIZE = ${SUBGROUP_SIZE}; -const uint NUM_SUBGROUPS = SG_GRID_X * SG_GRID_Y; -const uint WG_SIZE = NUM_SUBGROUPS * SUBGROUP_SIZE; - -const uint SG_TILE_M = WG_TILE_M / SG_GRID_Y; -const uint SG_TILE_N = WG_TILE_N / SG_GRID_X; -const uint MMAS_PER_SG_M = SG_TILE_M / MMA_M; -const uint MMAS_PER_SG_N = SG_TILE_N / MMA_N; - -const uint A_SLAB_INT8 = WG_TILE_M * MMA_K; -const uint B_USEFUL_U32 = MMA_K / 4u; -// No skew + coalesced write -- dbuf4zpg's B fix, unchanged (this file does -// not touch B staging at all). -const uint B_STRIDE_U32 = B_USEFUL_U32; -const uint B_SLAB_U32 = WG_TILE_N * B_STRIDE_U32; -const uint NUM_K_SLABS = WG_TILE_K / MMA_K; - -const uint A_SLAB_U32 = A_SLAB_INT8 / 4u; -const uint A_STRIDE_U32 = MMA_K / 4u; - -// One ping-pong slice covers all K-slabs of one chunk. -const uint ASH_SLICE_U32 = NUM_K_SLABS * A_SLAB_U32; -const uint BSH_SLICE_U32 = NUM_K_SLABS * B_SLAB_U32; - -// Double-buffered MMA operand staging. -shared uint Ash_int8[2u * ASH_SLICE_U32]; -shared uint Bsh_int8[2u * BSH_SLICE_U32]; - -// Per-WG-tile-row activation params (loaded ONCE at WG start; constant -// across groups). -shared int izp_sh[WG_TILE_M]; // int32 (cast from int8 source) for broadcast -shared float ifs_sh[WG_TILE_M]; // float32 (cast from fp16 source) for broadcast - -// Per-(group, output-channel) weight params, ping-ponged by group parity. -// (For per-channel INT8 only slice 0 is ever used.) -shared float wsc_sh[2u * WG_TILE_N]; -// SUM_g wsc[g][n]*wsum[g][n] per output channel -- weight-side only, so it is -// accumulated once in the prologue. dbuf4zpg's zp-hoist, unchanged. -shared float wcorr_sh[WG_TILE_N]; - -#ifdef HAS_BIAS -shared float bias_sh[WG_TILE_N]; -#endif - -#ifdef IO_TEXTURE -// Result staging for the imageStore epilogue, mirroring the fp16 kernel: -// SG_GRID_Y bands of MMA_M rows, each WG_TILE_N wide, row-major. A full -// WG_TILE_M x WG_TILE_N buffer would cost SG_GRID_Y/MMAS_PER_SG_M x more LDS -// and wreck occupancy. float16_t-typed because coopMatStore needs it. -const uint CSH_ROWS = SG_GRID_Y * MMA_M; -shared float16_t Csh_out[CSH_ROWS * WG_TILE_N]; -#endif - -// Running fp32 accumulator (across all groups). -coopmat - result[MMAS_PER_SG_M][MMAS_PER_SG_N]; - -// Per-group int32 MMA accumulator. -coopmat - accum_int32[MMAS_PER_SG_M][MMAS_PER_SG_N]; - - -// Byte-parallel int4 -> int8 widening. dbuf4zpg's, unchanged (B-side only). -// -// The four nibbles this shader needs from one packed uint are ALREADY one per -// byte (bits 3:0 of each byte for parity 0, bits 7:4 for parity 1), so all four -// can be widened at once instead of with a per-nibble -// shift/mask/bias-subtract/mask chain. -// -// For v in [0,15] the biased value is v-8. `v ^ 8` is exactly the 4-bit -// two's-complement pattern of v-8, because -8 == +8 (mod 16): -// v=0 -> 0x8 -> -8 v=7 -> 0xF -> -1 -// v=8 -> 0x0 -> 0 v=15 -> 0x7 -> +7 -// so the only remaining work is sign-extending bit 3 into bits 7:4 per byte. -// `sgn * 0x1E` does that with no cross-byte carry: 0x08 * 0x1E == 0xF0 exactly, -// and sgn is at most 0x08080808 so the product is at most 0xF0F0F0F0. -// -// A naive `nib - 0x08080808` would NOT work -- it borrows across byte lanes -// whenever a nibble is < 8. Shifts must be on uint, not int, so the >> is -// logical rather than arithmetic. -// -// ~5 ops per 4 weights vs ~22 for the per-nibble chain; bit-identical output. -uint widen_nibbles(const uint w, const uint parity) { - const uint nib = (parity == 0u) ? (w & 0x0F0F0F0Fu) : ((w >> 4u) & 0x0F0F0F0Fu); - const uint p = nib ^ 0x08080808u; - const uint sgn = p & 0x08080808u; - return p | (sgn * 0x1Eu); -} - -void main() { - const uvec2 tileID = uvec2(gl_WorkGroupID.xy); - const uvec2 warpInTile = uvec2( - gl_SubgroupID % SG_GRID_X, - gl_SubgroupID / SG_GRID_X); - - const uint K = uint(input_sizes.x); - const uint N = uint(output_sizes.x); - const uint N4 = (N + 3u) / 4u; - const uint nblocks_x_A = (K + 3u) >> 2u; - // A row stride in INT8 elements (dbuf4tr's binding is row-major int8, not - // the 4h4w ivec4 block layout dbuf4zpg used -- so A addressing below is in - // int8 elements, not int, and derived from nblocks_x_A so it matches the - // packer's `m_row * K4 + k4` addressing exactly (K % 4 == 0 makes them - // equal to K directly). - const uint a_row_stride_i8 = nblocks_x_A * 4u; - -#ifdef WEIGHT_INT4 - const uint num_groups = uint(num_groups_arg); - const uint CHUNKS_PER_GROUP = uint(K4_per_group) * 4u / WG_TILE_K; -#else - // Per-channel: a single quant "group" spanning all of K. The nested - // groups x chunks loop below collapses to a flat chunk loop, the wsum/wsc - // ping-pong never crosses a boundary, and the epilog runs exactly once. - const uint num_groups = 1u; - const uint CHUNKS_PER_GROUP = uint(num_groups_arg); -#endif - const uint num_chunks = num_groups * CHUNKS_PER_GROUP; - - const uint tile_m_start = WG_TILE_M * tileID.y; - const uint tile_n_start = WG_TILE_N * tileID.x; - - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - result[i][j] = coopmat(0.0); - accum_int32[i][j] = coopmat(0); - } - } - - const uint K_BLOCKS_PER_CHUNK = WG_TILE_K >> 2u; - - // --- A staging tile map (dbuf4tr's, unmodified): one MMA_M x MMA_K coopmat - // tile per subgroup per slot, dealt round-robin across the - // NUM_SUBGROUPS subgroups so every subgroup participates. Replaces - // dbuf4zpg's per-thread (m4, k4) map / a_active guard entirely -- see - // design.md D3 for why this is reused as-is, not re-derived. - const uint A_TILES_M = WG_TILE_M / MMA_M; - const uint A_TILES_K = WG_TILE_K / MMA_K; // == NUM_K_SLABS - const uint NUM_A_TILES = A_TILES_M * A_TILES_K; - const uint A_TILES_PER_SG = (NUM_A_TILES + NUM_SUBGROUPS - 1u) / NUM_SUBGROUPS; - -#ifdef WEIGHT_INT4 - // --- B staging thread map: (block, col) slots; each slot extracts one - // ColumnMajor LDS uint (4 K-contiguous sign-extended int8) --- - const uint B_TOTAL_SLOTS = K_BLOCKS_PER_CHUNK * WG_TILE_N; - const uint B_SLOTS_PER_THREAD = B_TOTAL_SLOTS / WG_SIZE; - const uint N8_PER_TILE = WG_TILE_N >> 3u; -#else - // --- B staging thread map: one (k4, n4) ivec4 block per active thread --- - const uint B_FETCH_SLOTS = K_BLOCKS_PER_CHUNK * (WG_TILE_N >> 2u); - const uint N4_PER_TILE = WG_TILE_N >> 2u; - const uint b_k4_in_chunk = gl_LocalInvocationID.x / N4_PER_TILE; - const uint b_n_uint_col = gl_LocalInvocationID.x % N4_PER_TILE; - const bool b_active = gl_LocalInvocationID.x < B_FETCH_SLOTS; -#endif - - // ===== INTERVENTION F: hoist loop-invariant B staging index math ===== - // dbuf4zpg's, unchanged -- B staging is untouched by this file's A-staging - // swap. See dbuf4zpg's header for the full rationale (ablation-attributed - // -16.8% block, mostly index arithmetic not memory traffic). -#ifdef WEIGHT_INT4 - // ===== bw: SOURCE-side change only -- one thread owns a whole (slab, col) - // column-slab and writes its 4 uints at CONSECUTIVE indices, so the four - // ds_write_b32 are merge-eligible. The shared array stays `uint` (the proven - // 4x packing ratio) and the compute-side coopMatLoad is byte-for-byte - // unchanged -- deliberately NOT the uvec4 retype, which fails correctness on - // this driver (16x packing; see the 8x coopMatStore precedent). - // comp/par depend only on the column, so a thread's 4 fetches share them and - // differ only in k4: 4 consecutive k4 blocks of the same n8 texel column. - const uint B_OWNERS = NUM_K_SLABS * WG_TILE_N; - const bool b_owns = gl_LocalInvocationID.x < B_OWNERS; - const uint b_slab = gl_LocalInvocationID.x / WG_TILE_N; - const uint b_col = gl_LocalInvocationID.x % WG_TILE_N; - const uint b_rem = b_col & 7u; - const uint b_comp1 = b_rem & 3u; - const uint b_par1 = b_rem >> 2u; - const uint b_n8blk1 = (tile_n_start >> 3u) + (b_col >> 3u); - const uint b_k4base = b_slab * (MMA_K >> 2u); - // same address the destination-oriented map produced, just all 4 in one thread - const uint b_lds_base = b_slab * B_SLAB_U32 + b_col * B_STRIDE_U32; -#endif - - // Prefetch temp registers. temp_A is a coopmat array (dbuf4tr's A-staging - // technique); indices into it are [[unroll]]-resolved compile-time - // constants, never dynamic -- dynamic indexing of a coopmat array is - // exactly the construct the Xclipse/AMD-PAL compiler has miscompiled - // before. - coopmat - temp_A[A_TILES_PER_SG]; -#ifdef WEIGHT_INT4 - ivec4 temp_B[MMA_K >> 2u]; // bw: one per k4 of the owned column-slab - float temp_wsc; -#else - ivec4 temp_B; -#endif - - // ========================================================= - // PROLOGUE - // ========================================================= - if (gl_LocalInvocationID.x < (WG_TILE_M >> 2u)) { - const uint m4 = (tile_m_start >> 2u) + gl_LocalInvocationID.x; - const vec4 sc = vec4(texelFetch(t_int8_input_scales, ivec3(m4, 0, 0), 0)); - const ivec4 zp = texelFetch(t_int8_input_zps, ivec3(m4, 0, 0), 0); - const uint base = gl_LocalInvocationID.x * 4u; - ifs_sh[base + 0u] = sc.x; ifs_sh[base + 1u] = sc.y; - ifs_sh[base + 2u] = sc.z; ifs_sh[base + 3u] = sc.w; - izp_sh[base + 0u] = zp.x; izp_sh[base + 1u] = zp.y; - izp_sh[base + 2u] = zp.z; izp_sh[base + 3u] = zp.w; - } - // Group 0 weight scales -> slice 0, and the hoisted weight-side correction - // SUM_g wsc[g][n]*wsum[g][n] accumulated across ALL groups. dbuf4zpg's - // zp-hoist, unchanged. - if (gl_LocalInvocationID.x < WG_TILE_N) { - const uint n_idx = tile_n_start + gl_LocalInvocationID.x; - f16vec4 sv0 = t_weight_scales[n_idx >> 2u]; - wsc_sh[gl_LocalInvocationID.x] = float(sv0[n_idx & 3u]); - - float corr = 0.0; - for (uint g = 0; g < num_groups; ++g) { - f16vec4 sv = t_weight_scales[g * N4 + (n_idx >> 2u)]; - corr += float(sv[n_idx & 3u]) * float(t_weight_sums[g * N + n_idx]); - } - wcorr_sh[gl_LocalInvocationID.x] = corr; - } - memoryBarrierShared(); - barrier(); - - // NOTE: dbuf4zpg builds izp_bcast/ifs_bcast AFTER the group loop, not here - // -- that is the register-pressure saving zp-hoist buys. Unchanged. - - // dbuf4: prefetch chunk 0 into temp registers, THEN store to slice 0 (no - // barrier here -- the main loop's first iteration barriers before - // reading slice 0). - // - // A staging (dbuf4tr's technique): per-subgroup coopMatLoad straight from - // the row-major global buffer. - [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { - const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; - if (t < NUM_A_TILES) { - const uint tm = t / A_TILES_K; - const uint tk = t % A_TILES_K; - coopMatLoad( - temp_A[s], t_packed_int8_input, - (tile_m_start + tm * MMA_M) * a_row_stride_i8 + tk * MMA_K, - a_row_stride_i8, - gl_CooperativeMatrixLayoutRowMajor); - } - } -#ifdef WEIGHT_INT4 - if (b_owns) { - [[unroll]] for (uint kk = 0; kk < (MMA_K >> 2u); ++kk) { -#ifdef WEIGHT_BUFFER - temp_B[kk] = t_packed_weight[(b_n8blk1 * nblocks_x_A) + b_k4base + kk]; -#else - temp_B[kk] = texelFetch(t_packed_weight, ivec2(b_k4base + kk, b_n8blk1), 0); -#endif - } - } -#else - if (b_active) { - const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; -#ifdef WEIGHT_BUFFER - temp_B = t_packed_weight[(b_k4_in_chunk * N4) + block_x_w]; -#else - temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, b_k4_in_chunk), 0); -#endif - } -#endif - { - // store chunk 0 -> slice 0 - // A staging (dbuf4tr's technique): coopMatStore into the same Ash_int8 - // slot layout dbuf4zpg's scalar scatter used to write. - [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { - const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; - if (t < NUM_A_TILES) { - const uint tm = t / A_TILES_K; - const uint tk = t % A_TILES_K; - coopMatStore( - temp_A[s], Ash_int8, - tk * A_SLAB_U32 + (tm * MMA_M) * A_STRIDE_U32, - A_STRIDE_U32, - gl_CooperativeMatrixLayoutRowMajor); - } - } -#ifdef WEIGHT_INT4 - if (b_owns) { - // 4 CONSECUTIVE uints from one thread -> merge-eligible into a wider - // store, versus 4 scalar stores from 4 different threads before. - [[unroll]] for (uint q = 0; q < (MMA_K >> 2u); ++q) { - const uint kk = (MMA_K >> 2u) - 1u - q; // DIAGNOSTIC: reverse order - Bsh_int8[b_lds_base + kk] = - widen_nibbles(uint(temp_B[kk][b_comp1]), b_par1); - } - } -#else - if (b_active) { - const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); - const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); - const uint n_col_base = b_n_uint_col * 4u; - [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { - Bsh_int8[slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab] = - uint(temp_B[n_in_blk]); - } - } -#endif - } - - // ========================================================= - // MAIN LOOP (dbuf4) — nested groups x chunks (kept nested; flattening it - // with a conditional coopmat epilog crashes the Xclipse PAL compiler at - // large spec-resolved trip counts). One barrier per chunk. Chunk - // iteration (global index `chunk`): - // 1. barrier — A/B slice (chunk%2) fully written; on the first chunk - // of group g, wsc slice (g%2) is too. - // 2. prefetch — chunk+1 (A tiles, B blocks) into temp; when chunk+1 - // starts a new group, also its wsc element. Skipped - // entirely on the final chunk. - // 3. int8 MMA — on slice (chunk%2) into accum_int32. - // 4. store — temp -> A/B slice ((chunk+1)%2), unpacking the weight; - // on a group boundary, wsc -> slice ((g+1)%2). - // The group epilog runs unconditionally at the tail of each group. - // ========================================================= - uint chunk = 0; - for (uint group_i = 0; group_i < num_groups; ++group_i) { - for (uint inner = 0; inner < CHUNKS_PER_GROUP; ++inner, ++chunk) { - const bool has_next = chunk + 1u < num_chunks; - const bool group_crossing = has_next && (inner + 1u == CHUNKS_PER_GROUP); - const uint cur_a = (chunk % 2u) * ASH_SLICE_U32; - const uint cur_b = (chunk % 2u) * BSH_SLICE_U32; - const uint nxt_a = ((chunk + 1u) % 2u) * ASH_SLICE_U32; - const uint nxt_b = ((chunk + 1u) % 2u) * BSH_SLICE_U32; - - // coopmat-lds-fence: barrier() alone does NOT order shared stores against a - // subsequent coopMatLoad on the M51 Xclipse/AMD-PAL driver -- symptom is one - // stale MMA_M-row band of A, all columns, ~2.5% of runs, no crash (observed - // 2026-09-02 in sdpa_compute_out_coopmat.glsl). Measured cost of the fence: - // none (see this change's results). See memory - // `coopmat-lds-needs-explicit-memorybarriershared`. - memoryBarrierShared(); - barrier(); - - // --- 2. prefetch chunk+1 -> temp --- - if (has_next) { - const uint chunkK_nxt = (chunk + 1u) * WG_TILE_K; - // A staging (dbuf4tr's technique): coopMatLoad straight from global. - [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { - const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; - if (t < NUM_A_TILES) { - const uint tm = t / A_TILES_K; - const uint tk = t % A_TILES_K; - coopMatLoad( - temp_A[s], t_packed_int8_input, - (tile_m_start + tm * MMA_M) * a_row_stride_i8 + chunkK_nxt + - tk * MMA_K, - a_row_stride_i8, - gl_CooperativeMatrixLayoutRowMajor); - } - } -#ifdef WEIGHT_INT4 - if (b_owns) { - [[unroll]] for (uint kk = 0; kk < (MMA_K >> 2u); ++kk) { - const uint k4_blk = (chunkK_nxt >> 2u) + b_k4base + kk; -#ifdef WEIGHT_BUFFER - temp_B[kk] = t_packed_weight[(b_n8blk1 * nblocks_x_A) + k4_blk]; -#else - temp_B[kk] = texelFetch(t_packed_weight, ivec2(k4_blk, b_n8blk1), 0); -#endif - } - } - if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { - const uint n_idx = tile_n_start + gl_LocalInvocationID.x; - f16vec4 sv = t_weight_scales[(group_i + 1u) * N4 + (n_idx >> 2u)]; - temp_wsc = float(sv[n_idx & 3u]); - } -#else - if (b_active) { - const uint block_y_w = (chunkK_nxt >> 2u) + b_k4_in_chunk; - const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; -#ifdef WEIGHT_BUFFER - temp_B = t_packed_weight[(block_y_w * N4) + block_x_w]; -#else - temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, block_y_w), 0); -#endif - } -#endif - } - - // --- 3. int8 MMA on the cur slice --- - [[unroll]] for (uint k = 0; k < NUM_K_SLABS; ++k) { - const uint slab_a_base_u32 = cur_a + k * A_SLAB_U32; - const uint slab_b_base_u32 = cur_b + k * B_SLAB_U32; - - coopmat matA[MMAS_PER_SG_M]; - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - const uint row_a = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); - coopMatLoad( - matA[i], Ash_int8, - slab_a_base_u32 + row_a * A_STRIDE_U32, - A_STRIDE_U32, - gl_CooperativeMatrixLayoutRowMajor); - } - - coopmat matB; - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint col_b = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopMatLoad( - matB, Bsh_int8, - slab_b_base_u32 + col_b * B_STRIDE_U32, - B_STRIDE_U32, - gl_CooperativeMatrixLayoutColumnMajor); - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - accum_int32[i][j] = coopMatMulAdd(matA[i], matB, accum_int32[i][j]); - } - } - } - - // --- 4. store temp (chunk+1) -> nxt slice --- - if (has_next) { - // A staging (dbuf4tr's technique): coopMatStore into the nxt slice. - [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { - const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; - if (t < NUM_A_TILES) { - const uint tm = t / A_TILES_K; - const uint tk = t % A_TILES_K; - coopMatStore( - temp_A[s], Ash_int8, - nxt_a + tk * A_SLAB_U32 + (tm * MMA_M) * A_STRIDE_U32, - A_STRIDE_U32, - gl_CooperativeMatrixLayoutRowMajor); - } - } -#ifdef WEIGHT_INT4 - if (b_owns) { - // 4 CONSECUTIVE uints from one thread -> merge-eligible into a wider - // store, versus 4 scalar stores from 4 different threads before. - [[unroll]] for (uint q = 0; q < (MMA_K >> 2u); ++q) { - const uint kk = (MMA_K >> 2u) - 1u - q; // DIAGNOSTIC: reverse order - Bsh_int8[nxt_b + b_lds_base + kk] = - widen_nibbles(uint(temp_B[kk][b_comp1]), b_par1); - } - } - if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { - const uint wbase_nxt = ((group_i + 1u) % 2u) * WG_TILE_N; - wsc_sh[wbase_nxt + gl_LocalInvocationID.x] = temp_wsc; - } -#else - if (b_active) { - const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); - const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); - const uint n_col_base = b_n_uint_col * 4u; - [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { - Bsh_int8[nxt_b + slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab] = - uint(temp_B[n_in_blk]); - } - } -#endif - } - } // chunks - - // --- Group epilog: scale-only accumulate, reset accum --- - // dbuf4zpg's, unchanged. Just result += float(acc) * wsc. The - // zero-point subtract and the ifs multiply are hoisted out of the group - // loop (applied once below). - { - const uint wbase = (group_i % 2u) * WG_TILE_N; - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - - coopmat wsc_bcast; - coopMatLoad( - wsc_bcast, wsc_sh, - wbase + local_n_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutRowMajor); - - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - result[i][j] += - coopmat( - accum_int32[i][j]) * wsc_bcast; - accum_int32[i][j] = coopmat(0); - } - } - } - } // groups - - // --- Hoisted correction, applied ONCE: -------------------------------- - // result = ifs * ( result - izp * SUM_g wsc_g*wsum_g ) - // dbuf4zpg's, unchanged. izp/ifs are loaded here rather than before the - // group loop so they are not live across it. - { - coopmat - izpf_bcast[MMAS_PER_SG_M]; - coopmat - ifs_bcast[MMAS_PER_SG_M]; - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - const uint local_m_base = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); - coopmat izp_i; - coopMatLoad( - izp_i, izp_sh, local_m_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutColumnMajor); - izpf_bcast[i] = - coopmat(izp_i); - coopMatLoad( - ifs_bcast[i], ifs_sh, local_m_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutColumnMajor); - } - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopmat wcorr_bcast; - coopMatLoad( - wcorr_bcast, wcorr_sh, local_n_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutRowMajor); - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - result[i][j] = ifs_bcast[i] * (result[i][j] - izpf_bcast[i] * wcorr_bcast); - } - } - } - - // --- Bias (optional) --- -#ifdef HAS_BIAS - if (apply_bias > 0) { - for (uint t = gl_LocalInvocationID.x; t < WG_TILE_N; t += WG_SIZE) { - bias_sh[t] = float(t_bias[tile_n_start + t]); - } - memoryBarrierShared(); - barrier(); - } -#endif - - // --- Store result tile --- - // N for the store address math MUST come from the spec constant, not the - // sizes UBO (see out_N_arg above). -#ifdef IO_TEXTURE - // Epilogue iteration i drains accumulator row-block i from EVERY subgroup - // into Csh_out at once, so the SG_GRID_Y bands it holds are disjoint global - // row ranges; the whole workgroup then imageStores them. lr / MMA_M is the - // writing subgroup's warpInTile.y, so the global row reproduces the buffer - // path's gi exactly. - // - // PORTABILITY NOTE: the barrier() in the loop body keeps this loop rolled - // despite [[unroll]], so result[i][j] IS dynamically indexed. Coopmat arrays - // are opaque per-lane storage and dynamic indexing is exactly the construct - // the Xclipse/AMD-PAL compiler has broken before -- check this first if the - // texture variants miscompile on M51. - const uint CSH_TEXELS_PER_ROW = WG_TILE_N / 4u; - const uint CSH_TEXELS = CSH_ROWS * CSH_TEXELS_PER_ROW; - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - // Guards Csh_out against the previous iteration's readers. Inert on i == 0 - // but must stay unconditional to remain workgroup-uniform. - // coopmat-lds-fence: barrier() alone does NOT order shared stores against a - // subsequent coopMatLoad on the M51 Xclipse/AMD-PAL driver -- symptom is one - // stale MMA_M-row band of A, all columns, ~2.5% of runs, no crash (observed - // 2026-09-02 in sdpa_compute_out_coopmat.glsl). Measured cost of the fence: - // none (see this change's results). See memory - // `coopmat-lds-needs-explicit-memorybarriershared`. - memoryBarrierShared(); - barrier(); - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { -#ifdef HAS_BIAS - if (apply_bias > 0) { - const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopmat bias_tile; - coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); - result[i][j] += bias_tile; - } -#endif - coopmat out_tile = - coopmat(result[i][j]); - coopMatStore( - out_tile, Csh_out, - warpInTile.y * MMA_M * WG_TILE_N + - MMA_N * (MMAS_PER_SG_N * warpInTile.x + j), - WG_TILE_N, - gl_CooperativeMatrixLayoutRowMajor); - } - memoryBarrierShared(); - barrier(); - - for (uint t = gl_LocalInvocationID.x; t < CSH_TEXELS; t += WG_SIZE) { - const uint lr = t / CSH_TEXELS_PER_ROW; - const uint lc4 = t % CSH_TEXELS_PER_ROW; - const uint m = - tile_m_start + (lr / MMA_M) * SG_TILE_M + i * MMA_M + (lr % MMA_M); - const uint base = lr * WG_TILE_N + lc4 * 4u; - imageStore( - t_output, - ivec3(tile_n_start / 4u + lc4, m, 0), - vec4( - float(Csh_out[base]), - float(Csh_out[base + 1u]), - float(Csh_out[base + 2u]), - float(Csh_out[base + 3u]))); - } - } -#else - const uint N_out = uint(out_N_arg); - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint gi = tile_m_start + MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); - const uint gj = tile_n_start + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - -#ifdef HAS_BIAS - if (apply_bias > 0) { - const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopmat bias_tile; - coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); - result[i][j] += bias_tile; - } -#endif - - coopmat out_tile = - coopmat(result[i][j]); - coopMatStore( - out_tile, t_output, - gi * N_out + gj, N_out, - gl_CooperativeMatrixLayoutRowMajor); - } - } -#endif // IO_TEXTURE -} diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbwr.yaml b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbwr.yaml deleted file mode 100644 index e7370ebafd2..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbwr.yaml +++ /dev/null @@ -1,68 +0,0 @@ -# !!! KNOWN-INCORRECT -- DO NOT PROMOTE, DO NOT TIME !!! -# Fails the correctness gate deterministically (3/3) on the num_groups==2 -# shapes only (K=256). Cause isolated 2026-09-08: the 4-wide-per-thread -# ownership with half the workgroup idle. Confirmed by dbuf4zpgbw2, which -# uses the SAME consecutive-write idea at 2-wide with all 256 threads -# active and passes 14/14 x3. Store ORDER is not the trigger (dbuf4zpgbwr, -# reversed, fails identically) and neither is the uvec4 retype (dbuf4zpgbw, -# no retype, fails identically). -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -# "zpg" + "-tr" combination: tsweep_dbuf4zpg with its per-thread scalar -# A-staging replaced by tsweep_dbuf4tr's coopMat-mediated A staging. B -# staging, zp-hoist, byte-parallel nibble widening are all dbuf4zpg's, -# unchanged -- only A-staging differs. Requires t_packed_int8_input in the -# ROW-MAJOR kPackedInt8_4W layout (same requirement as tsweep_dbuf4tr). -# -# PROMOTED 2026-09-01 as the shipped default -- see QuantizedLinear.cpp's -# dq8ca_coopmat_variant() for the full validation record (10/10 buffer, 6/6 -# texture3d correctness across 1B/3B/8B; +4.2% real e2e prefill; 46.50% -> -# 49.94% efficiency on 8B). Also selectable explicitly via -# ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4zpgbwr_txkgs. -# See openspec/changes/archive/2026-08-31-dq8ca-tr-staged-a-on-zpg. -# -# Single seed tile: the prior dbuf4zpg default's own tile (t128x64k32g42s32). -# A re-sweep against this shader's own (lower) register-pressure profile was -# run as a follow-up (coopmat-tr-tilesweep-4w-port) and found no better tile -# -- this remains the best known geometry. - -linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbwr: - parameter_names_with_default_values: - PRECISION: highp - HAS_BIAS: false - IO_STORAGE: buffer - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - MMA_M: 16 - MMA_N: 16 - MMA_K: 16 - WG_TILE_M: 64 - WG_TILE_N: 32 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - shader_variants: - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbwr_t128x64k32g42s32_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgbwr_t128x64k32g42s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr3.glsl b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr3.glsl deleted file mode 100644 index 24d42f12eff..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr3.glsl +++ /dev/null @@ -1,781 +0,0 @@ -/* - * Copyright (c) Meta Platforms, Inc. and affiliates. - * All rights reserved. - * - * This source code is licensed under the BSD-style license found in the - * LICENSE file in the root directory of this source tree. - */ - -/* - * "zpg" + "-tr" combination: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpg.glsl - * with its per-thread scalar A-staging replaced by - * linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4tr.glsl's coopMat-mediated A - * staging (coopMatLoad(global) -> coopmat<> -> coopMatStore(LDS)), PLUS tr3's - * shared-A layout: scalar int8_t storage, row-major over the FULL chunk, with - * A_ROW_PAD_I8 bytes of per-row padding. B staging, the quantization math and - * every correctness safeguard are dbuf4zpgtr's, byte-for-byte unchanged. This is - * an ADDITIVE combination, not a redesign: every non-A-staging block below - * (B staging: coalesced write, no skew; zp-hoist: izp/ifs applied once after - * the group loop via wcorr_sh; byte-parallel nibble widening; static - * A_ALWAYS_ACTIVE branch elision -- N/A here, see below; group epilog; - * bias/store epilogue) is byte-identical to dbuf4zpg's. Only the A-staging - * block (prologue load+store, main-loop prefetch+store) is dbuf4tr's, - * verbatim. - * - * dbuf4zpg's per-thread A staging used an `a_active` guard (statically always - * true when A_ACTIVE_THREADS == WG_SIZE, via the A_MAP_FULL-gated - * A_ALWAYS_ACTIVE macro). dbuf4tr's per-SUBGROUP tile map has no equivalent - * concept -- every subgroup participates via a `t < NUM_A_TILES` guard that - * depends only on gl_SubgroupID, not gl_LocalInvocationID.x -- so - * A_MAP_FULL/A_ALWAYS_ACTIVE is dropped entirely in this file; it would be - * dead code for the new A-staging block. - * - * Rationale for combining this way (not the reverse) and why this is worth - * building at all: see this change's design.md D0-D3. In short -- the only - * existing measurement of dbuf4tr's A-staging technique (28.72-30.51%, - * dq8ca-arch-redesign) was taken against dbuf4tr's own pre-zpg baseline - * (old B skew, no byte-parallel widening, no branch elision) -- a materially - * weaker shader than the 46.49-46.50% dbuf4zpg this file now combines it - * with. This file exists to answer whether that combination performs - * differently now that register pressure is already reduced. - * - * A staging (the actual delta from dbuf4zpg): - * dbuf4zpg: per-thread (m4, k4) ivec4 fetch, hoisted a_lds_off0/a_glb_row; - * only A_ACTIVE_THREADS invocations participate, each scattering - * 4 rows into shared A with 4 scalar stores. - * this file: per-SUBGROUP MMA_M x MMA_K tile fetch via coopMatLoad straight - * from a ROW-MAJOR (kPackedInt8_4W) int8 activation buffer, then - * coopMatStore into shared A -- dbuf4tr's mapping, - * unmodified (not re-derived; see design.md D3). - * - * t_packed_int8_input is therefore bound the same way dbuf4tr binds it: a - * SCALAR int8_t array in the kPackedInt8_4W layout (plain row-major int8, - * row stride K), produced by quantize_and_pack_4w_with_group_sums.glsl. - * QuantizedLinear.cpp's dq8ca_variant_wants_rowmajor_a() must recognize this - * file's variant token (tsweep_dbuf4zpgtr_t...) the same way it already - * recognizes tsweep_dbuf4tr_t/trm_t/trd_t, so graph-build time (packer - * selection) and dispatch time (kernel selection) cannot disagree. - * - * B CANNOT be coopmat-staged (int4 nibble unpack; a coopmat's per-lane layout - * is opaque to hand-assembly from unpacked registers) -- unchanged from both - * parent files. B staging below is dbuf4zpg's byte-parallel, coalesced, - * no-skew version, untouched. - * - * The loop structure is dbuf4's (both parents share it), unchanged: - * prologue: prefetch chunk 0 -> temp, store to slice 0 (no barrier) - * per iter: barrier -> prefetch(next) -> MMA(cur) -> store(next) - * kept nested (groups x chunks) with an unconditional group epilog -- - * flattening it crashes the Xclipse PAL compiler at large spec-resolved trip - * counts (see dbuf2's own header). - * - * Selected via - * ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4zpgtr_txkgs<32|64> - * (QuantizedLinear.cpp), additive to the tsweep_dbuf4zpg_t..., tsweep_dbuf4tr_t... - * and tsweep_t... namespaces. NOT the default -- unvalidated until it passes - * repeated test_llama_microbench --correctness-only runs (see - * dq8ca_coopmat_variant()'s comment on why a single pass is not proof). - * - * Performs: out[M,N] = dequant(int8_act) * dequant(int_w) (+ bias) - * via coopmat x coopmat -> coopmat on the matrix unit. - * - * Hard preconditions (dbuf4zpg's, plus dbuf4tr's row-major/alignment ones): - * M % WG_TILE_M == 0, N % WG_TILE_N == 0, K % WG_TILE_K == 0, - * group_size % WG_TILE_K == 0, K % 4 == 0, - * WG_TILE_M % MMA_M == 0, WG_TILE_K % MMA_K == 0, - * t_packed_int8_input in kPackedInt8_4W (row-major) layout, - * device exposes coopmatx-> at 16x16x16. - */ - -#version 450 core - -#extension GL_KHR_cooperative_matrix : require -#extension GL_KHR_memory_scope_semantics : require -#extension GL_KHR_shader_subgroup_basic : enable -#extension GL_EXT_shader_explicit_arithmetic_types : require -#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require -// 8-bit SSBO access: A is bound as a scalar int8_t array so that the -// coopMatLoad below has a MATCHING component type (see dbuf4tr's header for -// why the type must match on this driver). -#extension GL_EXT_shader_8bit_storage : require -#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require -#extension GL_EXT_control_flow_attributes : enable - -#define PRECISION ${PRECISION} - -$if WEIGHT_NBITS == 4: - #define WEIGHT_INT4 - -$if HAS_BIAS: - #define HAS_BIAS - -$if WEIGHT_STORAGE == "buffer": - #define WEIGHT_BUFFER - -$if IO_STORAGE == "texture3d": - #define IO_TEXTURE - -layout(std430) buffer; - -#include "common.glslh" - -// Bindings — match add_linear_dqa_qw_node arg order: -// output(0), fp_input(1), packed_int8_input(2), int_input_sums(3 - unused), -// input_scales(4), input_zps(5), packed_weight(6), weight_sums(7), -// weight_scales(8), bias(9). -${layout_declare_tensor(B, "w", "t_output", "half", IO_STORAGE, is_scalar_array=True)} -// t_input is unread here -- the activations arrive already quantized in -// t_packed_int8_input -- but stays declared so the binding layout matches the -// dispatch site. It tracks IO_STORAGE so the two IO tensors stay consistent. -${layout_declare_tensor(B, "r", "t_input", "half", IO_STORAGE, is_scalar_array=False)} -// ROW-MAJOR (kPackedInt8_4W) packed activations, bound as a scalar int8_t -// array (row stride = K int8) -- dbuf4tr's binding, unchanged. The stock -// 4h4w layout dbuf4zpg uses is NOT row-major (component index selects a row, -// non-affine), so it cannot be addressed by any coopMatLoad. -${layout_declare_tensor(B, "r", "t_packed_int8_input", "int8", "buffer", is_scalar_array=True)} -${layout_declare_tensor(B, "r", "t_int8_input_sums", "int", "buffer", is_scalar_array=True)} -${layout_declare_tensor(B, "r", "t_int8_input_scales", "half", "texture3d")} -${layout_declare_tensor(B, "r", "t_int8_input_zps", "int8", "texture3d")} -${layout_declare_tensor(B, "r", "t_packed_weight", "int", WEIGHT_STORAGE, is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_weight_sums", "int", "buffer", is_scalar_array=True)} -${layout_declare_tensor(B, "r", "t_weight_scales", "half", "buffer", is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_bias", "half", "buffer", is_scalar_array=True)} - -${layout_declare_ubo(B, "ivec4", "output_sizes")} -${layout_declare_ubo(B, "ivec4", "input_sizes")} - -layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; - -${layout_declare_spec_const(C, "int", "apply_bias", "0")} -// INT4 only; inert (0) for INT8 so the dispatcher's spec list lines up. -${layout_declare_spec_const(C, "int", "K4_per_group", "0")} -${layout_declare_spec_const(C, "int", "num_groups_arg", "0")} -${layout_declare_spec_const(C, "int", "out_N_arg", "0")} - -// Tile geometry -const uint MMA_M = ${MMA_M}; -const uint MMA_N = ${MMA_N}; -const uint MMA_K = ${MMA_K}; - -const uint WG_TILE_M = ${WG_TILE_M}; -const uint WG_TILE_N = ${WG_TILE_N}; -const uint WG_TILE_K = ${WG_TILE_K}; - -const uint SG_GRID_X = ${SG_GRID_X}; -const uint SG_GRID_Y = ${SG_GRID_Y}; -const uint SUBGROUP_SIZE = ${SUBGROUP_SIZE}; -const uint NUM_SUBGROUPS = SG_GRID_X * SG_GRID_Y; -const uint WG_SIZE = NUM_SUBGROUPS * SUBGROUP_SIZE; - -const uint SG_TILE_M = WG_TILE_M / SG_GRID_Y; -const uint SG_TILE_N = WG_TILE_N / SG_GRID_X; -const uint MMAS_PER_SG_M = SG_TILE_M / MMA_M; -const uint MMAS_PER_SG_N = SG_TILE_N / MMA_N; - -const uint B_USEFUL_U32 = MMA_K / 4u; -// No skew + coalesced write -- dbuf4zpg's B fix, unchanged (this file does -// not touch B staging at all). -const uint B_STRIDE_U32 = B_USEFUL_U32; -const uint B_SLAB_U32 = WG_TILE_N * B_STRIDE_U32; -const uint NUM_K_SLABS = WG_TILE_K / MMA_K; - -// ===== tr3-style A shared layout: THIS FILE'S ONLY STRUCTURAL CHANGE ===== -// dbuf4zpgtr stages A slab-major -- one MMA_K-wide slab per K step, row stride -// MMA_K/4 dwords, no padding, element type `uint` (4 int8 packed per element). -// tr3 instead keeps ONE row-major row per A row spanning the WHOLE chunk -// (WG_TILE_K wide), pads that row, and declares the array with the coopmat's -// own scalar element type so coopMatLoad/Store address it in elements. -// -// A_ROW_PAD_I8 pads the FULL-CHUNK row, NOT each MMA_K slab row -- the two have -// very different storage costs. Full-chunk padding costs WG_TILE_M*A_ROW_PAD_I8 -// bytes per slice; padding every slab row would cost NUM_K_SLABS times that. -// tr3's own value is 16 bytes (== ELEMENTS_PER_VEC4 for an 8-bit type). -const uint A_ROW_PAD_I8 = ${A_ROW_PAD_I8}u; -const uint STRIDE_A_I8 = WG_TILE_K + A_ROW_PAD_I8; -const uint ASH_SLICE_I8 = WG_TILE_M * STRIDE_A_I8; - -const uint BSH_SLICE_U32 = NUM_K_SLABS * B_SLAB_U32; - -// Double-buffered MMA operand staging. A is scalar int8_t (tr3-style); B keeps -// dbuf4zpg's packed-uint column-major slab layout untouched at this stage. -shared int8_t Ash_i8[2u * ASH_SLICE_I8]; -shared uint Bsh_int8[2u * BSH_SLICE_U32]; - -// Per-WG-tile-row activation params (loaded ONCE at WG start; constant -// across groups). -shared int izp_sh[WG_TILE_M]; // int32 (cast from int8 source) for broadcast -shared float ifs_sh[WG_TILE_M]; // float32 (cast from fp16 source) for broadcast - -// Per-(group, output-channel) weight params, ping-ponged by group parity. -// (For per-channel INT8 only slice 0 is ever used.) -shared float wsc_sh[2u * WG_TILE_N]; -// SUM_g wsc[g][n]*wsum[g][n] per output channel -- weight-side only, so it is -// accumulated once in the prologue. dbuf4zpg's zp-hoist, unchanged. -shared float wcorr_sh[WG_TILE_N]; - -#ifdef HAS_BIAS -shared float bias_sh[WG_TILE_N]; -#endif - -#ifdef IO_TEXTURE -// Result staging for the imageStore epilogue, mirroring the fp16 kernel: -// SG_GRID_Y bands of MMA_M rows, each WG_TILE_N wide, row-major. A full -// WG_TILE_M x WG_TILE_N buffer would cost SG_GRID_Y/MMAS_PER_SG_M x more LDS -// and wreck occupancy. float16_t-typed because coopMatStore needs it. -const uint CSH_ROWS = SG_GRID_Y * MMA_M; -shared float16_t Csh_out[CSH_ROWS * WG_TILE_N]; -#endif - -// Running fp32 accumulator (across all groups). -coopmat - result[MMAS_PER_SG_M][MMAS_PER_SG_N]; - -// Per-group int32 MMA accumulator. -coopmat - accum_int32[MMAS_PER_SG_M][MMAS_PER_SG_N]; - - -// Byte-parallel int4 -> int8 widening. dbuf4zpg's, unchanged (B-side only). -// -// The four nibbles this shader needs from one packed uint are ALREADY one per -// byte (bits 3:0 of each byte for parity 0, bits 7:4 for parity 1), so all four -// can be widened at once instead of with a per-nibble -// shift/mask/bias-subtract/mask chain. -// -// For v in [0,15] the biased value is v-8. `v ^ 8` is exactly the 4-bit -// two's-complement pattern of v-8, because -8 == +8 (mod 16): -// v=0 -> 0x8 -> -8 v=7 -> 0xF -> -1 -// v=8 -> 0x0 -> 0 v=15 -> 0x7 -> +7 -// so the only remaining work is sign-extending bit 3 into bits 7:4 per byte. -// `sgn * 0x1E` does that with no cross-byte carry: 0x08 * 0x1E == 0xF0 exactly, -// and sgn is at most 0x08080808 so the product is at most 0xF0F0F0F0. -// -// A naive `nib - 0x08080808` would NOT work -- it borrows across byte lanes -// whenever a nibble is < 8. Shifts must be on uint, not int, so the >> is -// logical rather than arithmetic. -// -// ~5 ops per 4 weights vs ~22 for the per-nibble chain; bit-identical output. -uint widen_nibbles(const uint w, const uint parity) { - const uint nib = (parity == 0u) ? (w & 0x0F0F0F0Fu) : ((w >> 4u) & 0x0F0F0F0Fu); - const uint p = nib ^ 0x08080808u; - const uint sgn = p & 0x08080808u; - return p | (sgn * 0x1Eu); -} - -void main() { - const uvec2 tileID = uvec2(gl_WorkGroupID.xy); - const uvec2 warpInTile = uvec2( - gl_SubgroupID % SG_GRID_X, - gl_SubgroupID / SG_GRID_X); - - const uint K = uint(input_sizes.x); - const uint N = uint(output_sizes.x); - const uint N4 = (N + 3u) / 4u; - const uint nblocks_x_A = (K + 3u) >> 2u; - // A row stride in INT8 elements (dbuf4tr's binding is row-major int8, not - // the 4h4w ivec4 block layout dbuf4zpg used -- so A addressing below is in - // int8 elements, not int, and derived from nblocks_x_A so it matches the - // packer's `m_row * K4 + k4` addressing exactly (K % 4 == 0 makes them - // equal to K directly). - const uint a_row_stride_i8 = nblocks_x_A * 4u; - -#ifdef WEIGHT_INT4 - const uint num_groups = uint(num_groups_arg); - const uint CHUNKS_PER_GROUP = uint(K4_per_group) * 4u / WG_TILE_K; -#else - // Per-channel: a single quant "group" spanning all of K. The nested - // groups x chunks loop below collapses to a flat chunk loop, the wsum/wsc - // ping-pong never crosses a boundary, and the epilog runs exactly once. - const uint num_groups = 1u; - const uint CHUNKS_PER_GROUP = uint(num_groups_arg); -#endif - const uint num_chunks = num_groups * CHUNKS_PER_GROUP; - - const uint tile_m_start = WG_TILE_M * tileID.y; - const uint tile_n_start = WG_TILE_N * tileID.x; - - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - result[i][j] = coopmat(0.0); - accum_int32[i][j] = coopmat(0); - } - } - - const uint K_BLOCKS_PER_CHUNK = WG_TILE_K >> 2u; - - // --- A staging tile map (dbuf4tr's, unmodified): one MMA_M x MMA_K coopmat - // tile per subgroup per slot, dealt round-robin across the - // NUM_SUBGROUPS subgroups so every subgroup participates. Replaces - // dbuf4zpg's per-thread (m4, k4) map / a_active guard entirely -- see - // design.md D3 for why this is reused as-is, not re-derived. - const uint A_TILES_M = WG_TILE_M / MMA_M; - const uint A_TILES_K = WG_TILE_K / MMA_K; // == NUM_K_SLABS - const uint NUM_A_TILES = A_TILES_M * A_TILES_K; - const uint A_TILES_PER_SG = (NUM_A_TILES + NUM_SUBGROUPS - 1u) / NUM_SUBGROUPS; - -#ifdef WEIGHT_INT4 - // --- B staging thread map: (block, col) slots; each slot extracts one - // ColumnMajor LDS uint (4 K-contiguous sign-extended int8) --- - const uint B_TOTAL_SLOTS = K_BLOCKS_PER_CHUNK * WG_TILE_N; - const uint B_SLOTS_PER_THREAD = B_TOTAL_SLOTS / WG_SIZE; - const uint N8_PER_TILE = WG_TILE_N >> 3u; -#else - // --- B staging thread map: one (k4, n4) ivec4 block per active thread --- - const uint B_FETCH_SLOTS = K_BLOCKS_PER_CHUNK * (WG_TILE_N >> 2u); - const uint N4_PER_TILE = WG_TILE_N >> 2u; - const uint b_k4_in_chunk = gl_LocalInvocationID.x / N4_PER_TILE; - const uint b_n_uint_col = gl_LocalInvocationID.x % N4_PER_TILE; - const bool b_active = gl_LocalInvocationID.x < B_FETCH_SLOTS; -#endif - - // ===== INTERVENTION F: hoist loop-invariant B staging index math ===== - // dbuf4zpg's, unchanged -- B staging is untouched by this file's A-staging - // swap. See dbuf4zpg's header for the full rationale (ablation-attributed - // -16.8% block, mostly index arithmetic not memory traffic). -#ifdef WEIGHT_INT4 - uint b_lds_off[B_SLOTS_PER_THREAD]; // LDS store offset within a slice - uint b_comp[B_SLOTS_PER_THREAD]; // which ivec4 component feeds this slot - uint b_par[B_SLOTS_PER_THREAD]; // nibble parity for this slot - uint b_n8blk[B_SLOTS_PER_THREAD]; // global texel column (N/8 blocks) - uint b_k4off[B_SLOTS_PER_THREAD]; // k4 offset of this slot within a chunk - [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { - const uint a = gl_LocalInvocationID.x + si * WG_SIZE; - const uint slab_idx = a / B_SLAB_U32; - const uint local_a = a % B_SLAB_U32; - const uint n_col = local_a / B_STRIDE_U32; - const uint k4_in_slab = local_a % B_STRIDE_U32; - const uint k4_in_chunk = slab_idx * (MMA_K >> 2u) + k4_in_slab; - const uint n8_in_tile = n_col >> 3u; - const uint rem = n_col & 7u; - b_lds_off[si] = a; - b_comp[si] = rem & 3u; - b_par[si] = rem >> 2u; - b_n8blk[si] = (tile_n_start >> 3u) + n8_in_tile; - b_k4off[si] = k4_in_chunk; - } -#endif - - // Prefetch temp registers. temp_A is a coopmat array (dbuf4tr's A-staging - // technique); indices into it are [[unroll]]-resolved compile-time - // constants, never dynamic -- dynamic indexing of a coopmat array is - // exactly the construct the Xclipse/AMD-PAL compiler has miscompiled - // before. - coopmat - temp_A[A_TILES_PER_SG]; -#ifdef WEIGHT_INT4 - ivec4 temp_B[B_SLOTS_PER_THREAD]; - float temp_wsc; -#else - ivec4 temp_B; -#endif - - // ========================================================= - // PROLOGUE - // ========================================================= - if (gl_LocalInvocationID.x < (WG_TILE_M >> 2u)) { - const uint m4 = (tile_m_start >> 2u) + gl_LocalInvocationID.x; - const vec4 sc = vec4(texelFetch(t_int8_input_scales, ivec3(m4, 0, 0), 0)); - const ivec4 zp = texelFetch(t_int8_input_zps, ivec3(m4, 0, 0), 0); - const uint base = gl_LocalInvocationID.x * 4u; - ifs_sh[base + 0u] = sc.x; ifs_sh[base + 1u] = sc.y; - ifs_sh[base + 2u] = sc.z; ifs_sh[base + 3u] = sc.w; - izp_sh[base + 0u] = zp.x; izp_sh[base + 1u] = zp.y; - izp_sh[base + 2u] = zp.z; izp_sh[base + 3u] = zp.w; - } - // Group 0 weight scales -> slice 0, and the hoisted weight-side correction - // SUM_g wsc[g][n]*wsum[g][n] accumulated across ALL groups. dbuf4zpg's - // zp-hoist, unchanged. - if (gl_LocalInvocationID.x < WG_TILE_N) { - const uint n_idx = tile_n_start + gl_LocalInvocationID.x; - f16vec4 sv0 = t_weight_scales[n_idx >> 2u]; - wsc_sh[gl_LocalInvocationID.x] = float(sv0[n_idx & 3u]); - - float corr = 0.0; - for (uint g = 0; g < num_groups; ++g) { - f16vec4 sv = t_weight_scales[g * N4 + (n_idx >> 2u)]; - corr += float(sv[n_idx & 3u]) * float(t_weight_sums[g * N + n_idx]); - } - wcorr_sh[gl_LocalInvocationID.x] = corr; - } - memoryBarrierShared(); - barrier(); - - // NOTE: dbuf4zpg builds izp_bcast/ifs_bcast AFTER the group loop, not here - // -- that is the register-pressure saving zp-hoist buys. Unchanged. - - // dbuf4: prefetch chunk 0 into temp registers, THEN store to slice 0 (no - // barrier here -- the main loop's first iteration barriers before - // reading slice 0). - // - // A staging (dbuf4tr's technique): per-subgroup coopMatLoad straight from - // the row-major global buffer. - [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { - const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; - if (t < NUM_A_TILES) { - const uint tm = t / A_TILES_K; - const uint tk = t % A_TILES_K; - coopMatLoad( - temp_A[s], t_packed_int8_input, - (tile_m_start + tm * MMA_M) * a_row_stride_i8 + tk * MMA_K, - a_row_stride_i8, - gl_CooperativeMatrixLayoutRowMajor); - } - } -#ifdef WEIGHT_INT4 - [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { -#ifdef WEIGHT_BUFFER - temp_B[si] = t_packed_weight[(b_n8blk[si] * nblocks_x_A) + b_k4off[si]]; -#else - temp_B[si] = texelFetch(t_packed_weight, ivec2(b_k4off[si], b_n8blk[si]), 0); -#endif - } -#else - if (b_active) { - const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; -#ifdef WEIGHT_BUFFER - temp_B = t_packed_weight[(b_k4_in_chunk * N4) + block_x_w]; -#else - temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, b_k4_in_chunk), 0); -#endif - } -#endif - { - // store chunk 0 -> slice 0 - // A staging (dbuf4tr's technique): coopMatStore into the same Ash_int8 - // slot layout dbuf4zpg's scalar scatter used to write. - [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { - const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; - if (t < NUM_A_TILES) { - const uint tm = t / A_TILES_K; - const uint tk = t % A_TILES_K; - coopMatStore( - temp_A[s], Ash_i8, - (tm * MMA_M) * STRIDE_A_I8 + tk * MMA_K, - STRIDE_A_I8, - gl_CooperativeMatrixLayoutRowMajor); - } - } -#ifdef WEIGHT_INT4 - [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { - Bsh_int8[b_lds_off[si]] = - widen_nibbles(uint(temp_B[si][b_comp[si]]), b_par[si]); - } -#else - if (b_active) { - const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); - const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); - const uint n_col_base = b_n_uint_col * 4u; - [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { - Bsh_int8[slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab] = - uint(temp_B[n_in_blk]); - } - } -#endif - } - - // ========================================================= - // MAIN LOOP (dbuf4) — nested groups x chunks (kept nested; flattening it - // with a conditional coopmat epilog crashes the Xclipse PAL compiler at - // large spec-resolved trip counts). One barrier per chunk. Chunk - // iteration (global index `chunk`): - // 1. barrier — A/B slice (chunk%2) fully written; on the first chunk - // of group g, wsc slice (g%2) is too. - // 2. prefetch — chunk+1 (A tiles, B blocks) into temp; when chunk+1 - // starts a new group, also its wsc element. Skipped - // entirely on the final chunk. - // 3. int8 MMA — on slice (chunk%2) into accum_int32. - // 4. store — temp -> A/B slice ((chunk+1)%2), unpacking the weight; - // on a group boundary, wsc -> slice ((g+1)%2). - // The group epilog runs unconditionally at the tail of each group. - // ========================================================= - uint chunk = 0; - for (uint group_i = 0; group_i < num_groups; ++group_i) { - for (uint inner = 0; inner < CHUNKS_PER_GROUP; ++inner, ++chunk) { - const bool has_next = chunk + 1u < num_chunks; - const bool group_crossing = has_next && (inner + 1u == CHUNKS_PER_GROUP); - const uint cur_a = (chunk % 2u) * ASH_SLICE_I8; - const uint cur_b = (chunk % 2u) * BSH_SLICE_U32; - const uint nxt_a = ((chunk + 1u) % 2u) * ASH_SLICE_I8; - const uint nxt_b = ((chunk + 1u) % 2u) * BSH_SLICE_U32; - - // coopmat-lds-fence: barrier() alone does NOT order shared stores against a - // subsequent coopMatLoad on the M51 Xclipse/AMD-PAL driver -- symptom is one - // stale MMA_M-row band of A, all columns, ~2.5% of runs, no crash (observed - // 2026-09-02 in sdpa_compute_out_coopmat.glsl). Measured cost of the fence: - // none (see this change's results). See memory - // `coopmat-lds-needs-explicit-memorybarriershared`. - memoryBarrierShared(); - barrier(); - - // --- 2. prefetch chunk+1 -> temp --- - if (has_next) { - const uint chunkK_nxt = (chunk + 1u) * WG_TILE_K; - // A staging (dbuf4tr's technique): coopMatLoad straight from global. - [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { - const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; - if (t < NUM_A_TILES) { - const uint tm = t / A_TILES_K; - const uint tk = t % A_TILES_K; - coopMatLoad( - temp_A[s], t_packed_int8_input, - (tile_m_start + tm * MMA_M) * a_row_stride_i8 + chunkK_nxt + - tk * MMA_K, - a_row_stride_i8, - gl_CooperativeMatrixLayoutRowMajor); - } - } -#ifdef WEIGHT_INT4 - [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { - const uint k4_blk = (chunkK_nxt >> 2u) + b_k4off[si]; -#ifdef WEIGHT_BUFFER - temp_B[si] = t_packed_weight[(b_n8blk[si] * nblocks_x_A) + k4_blk]; -#else - temp_B[si] = texelFetch(t_packed_weight, ivec2(k4_blk, b_n8blk[si]), 0); -#endif - } - if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { - const uint n_idx = tile_n_start + gl_LocalInvocationID.x; - f16vec4 sv = t_weight_scales[(group_i + 1u) * N4 + (n_idx >> 2u)]; - temp_wsc = float(sv[n_idx & 3u]); - } -#else - if (b_active) { - const uint block_y_w = (chunkK_nxt >> 2u) + b_k4_in_chunk; - const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; -#ifdef WEIGHT_BUFFER - temp_B = t_packed_weight[(block_y_w * N4) + block_x_w]; -#else - temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, block_y_w), 0); -#endif - } -#endif - } - - // --- 3. int8 MMA on the cur slice --- - [[unroll]] for (uint k = 0; k < NUM_K_SLABS; ++k) { - // row-major full-chunk A: the K step is a column offset within the - // row, not a separate slab base. - const uint k_col_a_i8 = k * MMA_K; - const uint slab_b_base_u32 = cur_b + k * B_SLAB_U32; - - coopmat matA[MMAS_PER_SG_M]; - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - const uint row_a = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); - coopMatLoad( - matA[i], Ash_i8, - cur_a + row_a * STRIDE_A_I8 + k_col_a_i8, - STRIDE_A_I8, - gl_CooperativeMatrixLayoutRowMajor); - } - - coopmat matB; - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint col_b = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopMatLoad( - matB, Bsh_int8, - slab_b_base_u32 + col_b * B_STRIDE_U32, - B_STRIDE_U32, - gl_CooperativeMatrixLayoutColumnMajor); - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - accum_int32[i][j] = coopMatMulAdd(matA[i], matB, accum_int32[i][j]); - } - } - } - - // --- 4. store temp (chunk+1) -> nxt slice --- - if (has_next) { - // A staging (dbuf4tr's technique): coopMatStore into the nxt slice. - [[unroll]] for (uint s = 0; s < A_TILES_PER_SG; ++s) { - const uint t = gl_SubgroupID + s * NUM_SUBGROUPS; - if (t < NUM_A_TILES) { - const uint tm = t / A_TILES_K; - const uint tk = t % A_TILES_K; - coopMatStore( - temp_A[s], Ash_i8, - nxt_a + (tm * MMA_M) * STRIDE_A_I8 + tk * MMA_K, - STRIDE_A_I8, - gl_CooperativeMatrixLayoutRowMajor); - } - } -#ifdef WEIGHT_INT4 - [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { - Bsh_int8[nxt_b + b_lds_off[si]] = - widen_nibbles(uint(temp_B[si][b_comp[si]]), b_par[si]); - } - if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { - const uint wbase_nxt = ((group_i + 1u) % 2u) * WG_TILE_N; - wsc_sh[wbase_nxt + gl_LocalInvocationID.x] = temp_wsc; - } -#else - if (b_active) { - const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); - const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); - const uint n_col_base = b_n_uint_col * 4u; - [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { - Bsh_int8[nxt_b + slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab] = - uint(temp_B[n_in_blk]); - } - } -#endif - } - } // chunks - - // --- Group epilog: scale-only accumulate, reset accum --- - // dbuf4zpg's, unchanged. Just result += float(acc) * wsc. The - // zero-point subtract and the ifs multiply are hoisted out of the group - // loop (applied once below). - { - const uint wbase = (group_i % 2u) * WG_TILE_N; - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - - coopmat wsc_bcast; - coopMatLoad( - wsc_bcast, wsc_sh, - wbase + local_n_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutRowMajor); - - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - result[i][j] += - coopmat( - accum_int32[i][j]) * wsc_bcast; - accum_int32[i][j] = coopmat(0); - } - } - } - } // groups - - // --- Hoisted correction, applied ONCE: -------------------------------- - // result = ifs * ( result - izp * SUM_g wsc_g*wsum_g ) - // dbuf4zpg's, unchanged. izp/ifs are loaded here rather than before the - // group loop so they are not live across it. - { - coopmat - izpf_bcast[MMAS_PER_SG_M]; - coopmat - ifs_bcast[MMAS_PER_SG_M]; - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - const uint local_m_base = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); - coopmat izp_i; - coopMatLoad( - izp_i, izp_sh, local_m_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutColumnMajor); - izpf_bcast[i] = - coopmat(izp_i); - coopMatLoad( - ifs_bcast[i], ifs_sh, local_m_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutColumnMajor); - } - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopmat wcorr_bcast; - coopMatLoad( - wcorr_bcast, wcorr_sh, local_n_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutRowMajor); - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - result[i][j] = ifs_bcast[i] * (result[i][j] - izpf_bcast[i] * wcorr_bcast); - } - } - } - - // --- Bias (optional) --- -#ifdef HAS_BIAS - if (apply_bias > 0) { - for (uint t = gl_LocalInvocationID.x; t < WG_TILE_N; t += WG_SIZE) { - bias_sh[t] = float(t_bias[tile_n_start + t]); - } - memoryBarrierShared(); - barrier(); - } -#endif - - // --- Store result tile --- - // N for the store address math MUST come from the spec constant, not the - // sizes UBO (see out_N_arg above). -#ifdef IO_TEXTURE - // Epilogue iteration i drains accumulator row-block i from EVERY subgroup - // into Csh_out at once, so the SG_GRID_Y bands it holds are disjoint global - // row ranges; the whole workgroup then imageStores them. lr / MMA_M is the - // writing subgroup's warpInTile.y, so the global row reproduces the buffer - // path's gi exactly. - // - // PORTABILITY NOTE: the barrier() in the loop body keeps this loop rolled - // despite [[unroll]], so result[i][j] IS dynamically indexed. Coopmat arrays - // are opaque per-lane storage and dynamic indexing is exactly the construct - // the Xclipse/AMD-PAL compiler has broken before -- check this first if the - // texture variants miscompile on M51. - const uint CSH_TEXELS_PER_ROW = WG_TILE_N / 4u; - const uint CSH_TEXELS = CSH_ROWS * CSH_TEXELS_PER_ROW; - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - // Guards Csh_out against the previous iteration's readers. Inert on i == 0 - // but must stay unconditional to remain workgroup-uniform. - // coopmat-lds-fence: barrier() alone does NOT order shared stores against a - // subsequent coopMatLoad on the M51 Xclipse/AMD-PAL driver -- symptom is one - // stale MMA_M-row band of A, all columns, ~2.5% of runs, no crash (observed - // 2026-09-02 in sdpa_compute_out_coopmat.glsl). Measured cost of the fence: - // none (see this change's results). See memory - // `coopmat-lds-needs-explicit-memorybarriershared`. - memoryBarrierShared(); - barrier(); - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { -#ifdef HAS_BIAS - if (apply_bias > 0) { - const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopmat bias_tile; - coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); - result[i][j] += bias_tile; - } -#endif - coopmat out_tile = - coopmat(result[i][j]); - coopMatStore( - out_tile, Csh_out, - warpInTile.y * MMA_M * WG_TILE_N + - MMA_N * (MMAS_PER_SG_N * warpInTile.x + j), - WG_TILE_N, - gl_CooperativeMatrixLayoutRowMajor); - } - memoryBarrierShared(); - barrier(); - - for (uint t = gl_LocalInvocationID.x; t < CSH_TEXELS; t += WG_SIZE) { - const uint lr = t / CSH_TEXELS_PER_ROW; - const uint lc4 = t % CSH_TEXELS_PER_ROW; - const uint m = - tile_m_start + (lr / MMA_M) * SG_TILE_M + i * MMA_M + (lr % MMA_M); - const uint base = lr * WG_TILE_N + lc4 * 4u; - imageStore( - t_output, - ivec3(tile_n_start / 4u + lc4, m, 0), - vec4( - float(Csh_out[base]), - float(Csh_out[base + 1u]), - float(Csh_out[base + 2u]), - float(Csh_out[base + 3u]))); - } - } -#else - const uint N_out = uint(out_N_arg); - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint gi = tile_m_start + MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); - const uint gj = tile_n_start + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - -#ifdef HAS_BIAS - if (apply_bias > 0) { - const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopmat bias_tile; - coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); - result[i][j] += bias_tile; - } -#endif - - coopmat out_tile = - coopmat(result[i][j]); - coopMatStore( - out_tile, t_output, - gi * N_out + gj, N_out, - gl_CooperativeMatrixLayoutRowMajor); - } - } -#endif // IO_TEXTURE -} diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr3.yaml b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr3.yaml deleted file mode 100644 index c4665a771d4..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr3.yaml +++ /dev/null @@ -1,92 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -# "zpgtr3": dbuf4zpgtr with tr3's SHARED-A layout. Three changes vs zpgtr, all -# on the A path only: -# 1. shared storage element type uint (4x int8 packed) -> scalar int8_t -# 2. slab-major (one MMA_K-wide slab per K step) -> row-major over the FULL -# chunk (WG_TILE_K wide) -# 3. A_ROW_PAD_I8 bytes of padding on each FULL-CHUNK row (tr3 uses 16) -# B staging, zp-hoist, byte-parallel nibble widening, the group/chunk nesting, -# the memoryBarrierShared fences and the fp16 output path are dbuf4zpgtr's, -# byte-for-byte unchanged. -# -# OPT-IN ONLY -- not a default. Select with -# ET_VK_DQ8CA_COOPMAT_VARIANT=tsweep_dbuf4zpgtr3_txkgsp -# -# LDS cost of the A array is 2*WG_TILE_M*(WG_TILE_K+A_ROW_PAD_I8) bytes, i.e. -# padding is paid once per row, NOT once per MMA_K slab row. At -# t128x64k32g42s32 that is p0 8192 B / p8 10240 B / p16 12288 B against -# zpgtr's 8192 B, so p16 adds 4 KiB and is the only one of the three that can -# cost a workgroup of LDS-limited occupancy. - -linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr3: - parameter_names_with_default_values: - PRECISION: highp - HAS_BIAS: false - IO_STORAGE: buffer - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - MMA_M: 16 - MMA_N: 16 - MMA_K: 16 - WG_TILE_M: 64 - WG_TILE_N: 32 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - A_ROW_PAD_I8: 16 - shader_variants: - # --- shipped geometry, pad swept: isolates padding from layout+element type - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr3_t128x64k32g42s32p16_buffer_texture2d_half - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - A_ROW_PAD_I8: 16 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr3_t128x64k32g42s32p8_buffer_texture2d_half - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - A_ROW_PAD_I8: 8 - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr3_t128x64k32g42s32p0_buffer_texture2d_half - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - A_ROW_PAD_I8: 0 - # --- iso geometry (the config TR3 was measured at) with tr3's own pad - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr3_t128x128k32g42s32p16_buffer_texture2d_half - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 32 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - A_ROW_PAD_I8: 16 - # --- texture-IO coverage at the shipped geometry (storage-variant parity) - - NAME: linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr3_t128x64k32g42s32p16_texture3d_texture2d_half - IO_STORAGE: texture3d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - A_ROW_PAD_I8: 16 - # NOTE on the INT8 weight path: the baseline dbuf4zpgtr yaml ships no - # WEIGHT_NBITS=8 variant either, and the host builds kernel names as - # ___half with no nbits field, so a w8 entry would be - # unreachable. The `#else` (INT8) branch of the shader is untouched by this - # change; it is simply not compiled, exactly as in the baseline. diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_coop.glsl b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_coop.glsl deleted file mode 100644 index b7f14f435fa..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_coop.glsl +++ /dev/null @@ -1,192 +0,0 @@ -/* - * Copyright (c) Meta Platforms, Inc. and affiliates. - * All rights reserved. - * - * This source code is licensed under the BSD-style license found in the - * LICENSE file in the root directory of this source tree. - */ - -#version 450 core - -#define PRECISION ${PRECISION} -#define VEC4_T ${texel_load_type(DTYPE, IO_STORAGE)} -#define T ${texel_load_component_type(DTYPE, IO_STORAGE)} - -$if IO_STORAGE == "buffer": - #define OUTPUT_BUFFER - #define INPUT_BUFFER - #define ATTN_WEIGHTS_BUFFER -$if K_CACHE_STORAGE == "buffer": - #define K_CACHE_BUFFER - -#define Q_LAYOUT DHSB -#define K_LAYOUT DHSB - -#define TILE_K4 ${TILE_K4} -#define TILE_N4 ${TILE_N4} - -#define TILE_M 1 -#define TILE_K ${TILE_K4 * 4} -#define TILE_N ${TILE_N4 * 4} - -#define NUM_WORKERS_PER_OUT 64 - -${define_required_extensions(IO_STORAGE, DTYPE)} - -layout(std430) buffer; - -#include "common.glslh" - -${layout_declare_tensor(B, "w", "t_attn_weights", DTYPE, IO_STORAGE, is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_q", DTYPE, IO_STORAGE, is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_k", DTYPE, K_CACHE_STORAGE, is_scalar_array=False)} - -${layout_declare_ubo(B, "ivec4", "q_sizes")} -${layout_declare_ubo(B, "ivec4", "k_sizes")} -${layout_declare_ubo(B, "int", "input_pos")} - -layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; - -${layout_declare_spec_const(C, "float", "inv_scale", "1.0")} - -#include "sdpa_fp_q_projected_tile_load.glslh" -#include "sdpa_fp_k_cache_tile_load.glslh" -#include "linear_fp_output_tile_fp_compute.glslh" -#include "sdpa_fp_attn_weight_tile_store.glslh" - -shared FPOutTile partial_sums[NUM_WORKERS_PER_OUT]; - -/* - * See the tiled variant of this shader for the implemented behavior. This - * shader is implements an optimization for cases where sequence length is 1; in - * these cases, the matrix multiplication being performed is akin to gemv, which - * benefits from using a co-operative algorithm for reduction. For this shader - * the entire work group co-operates to compute one reduction output. - */ - -void main() { - const int worker_id = int(gl_LocalInvocationID.y); - - const int tile_idx_x = int(gl_GlobalInvocationID.x); - // idx along output num_q_heads dim - const int q_h = int(gl_GlobalInvocationID.z); - - // idx along the output context_len dim - const int c = tile_idx_x * TILE_N; - const int c4 = div_4(c); - - // idx along the output seq_len dim. Note that for this shader seq_len will be - // 1. - const int s = 0; - - // head dimension - const int D = q_sizes.x; - // texel size of head_dim, over which the dot product is accumulated - const int D4 = div_up_4(D); - // number of Q heads - const int Q_H = q_sizes.y; - // sequence length - const int S = q_sizes.z; - const int S_aligned = align_up_4(S); - - // number of K/V heads - const int KV_H = k_sizes.y; - // Max context length - const int C = k_sizes.z; - const int C4 = div_up_4(C); - - int kv_h = q_h; - if (KV_H < Q_H) { - kv_h = q_h / (Q_H / KV_H); - } - - // current context length - const int context_len = input_pos + S; - const int context_texel_len = div_up_4(context_len); - - // bounds check - if (c >= context_len || s >= S || q_h >= Q_H) { - return; - } - - FPOutTile out_tile; - initialize(out_tile); - - FPInputTile q_tile; - FPWeightTile w_tile; - - // If the tile is completely inside the mask region, then there is no need to - // compute the output tile. All the elements in the output tile can be set to - // negative infinity. - bool tile_in_mask_region = c > (input_pos + s + (TILE_M - 1)); - if (tile_in_mask_region) { - const VEC4_T negative_infinity_vec = VEC4_T(negative_infinity_val); - set_out_tile_to_vec(out_tile, negative_infinity_vec); - } - // Otherwise, need to actually compute output tile - else { - for (int d4 = worker_id; d4 < D4; d4 += NUM_WORKERS_PER_OUT) { - load_q_projected_tile_with_checks( - q_tile, - d4, - s, - q_h, - D4, - D, - S, - Q_H); - - load_k_cache_tile_with_checks( - w_tile, - d4, - c, - kv_h, - D4, - D, - context_len, - C, - KV_H); - - fp_accumulate_with_fp_weight(out_tile, q_tile, w_tile); - } - } - - partial_sums[worker_id] = out_tile; - - memoryBarrierShared(); - barrier(); - - // Tree reduction to compute the overall result. - for (int i = NUM_WORKERS_PER_OUT / 2; i > 0; i /= 2) { - if (worker_id < i) { - accumulate_out_tile_with_out_tile( - partial_sums[worker_id], partial_sums[worker_id + i]); - } - memoryBarrierShared(); - barrier(); - } - - // Only the first thread will write out the result - if (worker_id == 0) { - out_tile = partial_sums[0]; - // Apply scale and mask if the tile was not entirely in the mask region - if (!tile_in_mask_region) { - VEC4_T inv_scale_vec = VEC4_T(inv_scale); - apply_scale_and_mask( - out_tile, - inv_scale_vec, - input_pos, - c, - s); - } - - store_attn_weight_tile_with_checks( - out_tile, - c4, - s, - q_h, - context_texel_len, - S_aligned, - Q_H); - } -} diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_coop.yaml b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_coop.yaml deleted file mode 100644 index d5cadc36060..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_coop.yaml +++ /dev/null @@ -1,25 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -sdpa_compute_attn_weights_coop: - parameter_names_with_default_values: - DTYPE: float - IO_STORAGE: texture3d - K_CACHE_STORAGE: texture3d - TILE_K4: 1 - TILE_N4: 1 - generate_variant_forall: - combination: - parameter_names: [IO_STORAGE, K_CACHE_STORAGE] - combos: - - parameter_values: [texture3d, texture3d] - - parameter_values: [buffer, texture3d] - - parameter_values: [buffer, buffer] - DTYPE: - - VALUE: float - - VALUE: half - shader_variants: - - NAME: sdpa_compute_attn_weights_coop diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coop.glsl b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coop.glsl deleted file mode 100644 index cd2c689ebc8..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coop.glsl +++ /dev/null @@ -1,199 +0,0 @@ -/* - * Copyright (c) Meta Platforms, Inc. and affiliates. - * All rights reserved. - * - * This source code is licensed under the BSD-style license found in the - * LICENSE file in the root directory of this source tree. - */ - -#version 450 core - -#define PRECISION ${PRECISION} -#define VEC4_T ${texel_load_type(DTYPE, IO_STORAGE)} -#define T ${texel_load_component_type(DTYPE, IO_STORAGE)} - -$if IO_STORAGE == "buffer": - #define OUTPUT_BUFFER - #define INPUT_BUFFER - #define ATTN_WEIGHTS_BUFFER -$if V_CACHE_STORAGE == "buffer": - #define V_CACHE_BUFFER - -#define V_LAYOUT DHSB -#define OUT_LAYOUT DHSB -#define SDPA_V_BUF t_v_cache - -#define TILE_K4 ${TILE_K4} -#define TILE_N4 ${TILE_N4} - -#define TILE_M 1 -#define TILE_K ${TILE_K4 * 4} -#define TILE_N ${TILE_N4 * 4} - -#define NUM_WORKERS_PER_OUT 64 - -${define_required_extensions(IO_STORAGE, DTYPE)} - -layout(std430) buffer; - -#include "common.glslh" - -${layout_declare_tensor(B, "w", "t_output", DTYPE, IO_STORAGE, is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_attn_weights", DTYPE, IO_STORAGE, is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_v_cache", DTYPE, V_CACHE_STORAGE, is_scalar_array=False)} - -${layout_declare_ubo(B, "ivec4", "q_sizes")} -${layout_declare_ubo(B, "ivec4", "v_sizes")} -${layout_declare_ubo(B, "int", "input_pos")} - -layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; - -${layout_declare_spec_const(C, "float", "inv_scale", "1.0")} - -#include "sdpa_fp_attn_weight_tile_load.glslh" -#include "sdpa_fp_v_cache_tile_load.glslh" -#include "linear_fp_output_tile_fp_compute.glslh" -#include "sdpa_fp_out_tile_store.glslh" - -shared FPOutTile partial_sums[NUM_WORKERS_PER_OUT]; - -/* - * See the tiled variant of this shader for the implemented behavior. This - * shader is implements an optimization for cases where sequence length is 1; in - * these cases, the matrix multiplication being performed is akin to gemv, which - * benefits from using a co-operative algorithm for reduction. For this shader - * the entire work group co-operates to compute one reduction output. - */ - -void main() { - const int worker_id = int(gl_LocalInvocationID.y); - - const int tile_idx_x = int(gl_GlobalInvocationID.x); - // idx along output num_q_heads dim - const int q_h = int(gl_GlobalInvocationID.z); - - // idx along the output head_dim dim - const int d = tile_idx_x * TILE_N; - const int d4 = div_4(d); - - // idx along the output seq_len dim. Note that for this shader seq_len will be - // 1. - const int s = 0; - - // texel size of head_dim - const int D4 = div_up_4(q_sizes.x); - // number of Q heads - const int Q_H = q_sizes.y; - // sequence length - const int S = q_sizes.z; - const int S_aligned = align_up_4(S); - - // number of K/V heads - const int KV_H = v_sizes.y; - // Max context length - const int C = v_sizes.z; - const int C4 = div_up_4(C); - - int kv_h = q_h; - if (KV_H < Q_H) { - kv_h = q_h / (Q_H / KV_H); - } - - // current context length - const int context_len = input_pos + S; - const int context_texel_len = div_up_4(context_len); - - // bounds check - if (d4 >= D4 || s >= S || q_h >= Q_H) { - return; - } - - FPOutTile out_tile; - initialize(out_tile); - - FPInputTile attn_weight_tile; - FPWeightTile w_tile; - - const int context_len_aligned_down = context_len - mod_4(context_len); - const int C4_limit = div_up_4(context_len_aligned_down); - - for (int c4 = worker_id; c4 < C4_limit; c4 += NUM_WORKERS_PER_OUT) { - const int c = mul_4(c4); - - load_attn_weight_tile_no_checks( - attn_weight_tile, - c4, - s, - q_h, - context_texel_len, - S_aligned, - Q_H); - - load_v_cache_tile_no_checks( - w_tile, - d4, - c, - kv_h, - D4, - context_len, - C, - KV_H); - - fp_accumulate_with_fp_weight(out_tile, attn_weight_tile, w_tile); - } - // first worker in the work group will handle final texel, which may contain - // padding elements. - if (worker_id == 0) { - for (int c4 = C4_limit; c4 < context_texel_len; c4++) { - const int c = mul_4(c4); - load_attn_weight_tile_with_checks( - attn_weight_tile, - c4, - s, - q_h, - context_texel_len, - S_aligned, - Q_H); - - load_v_cache_tile_with_checks( - w_tile, - d4, - c, - kv_h, - D4, - context_len, - C, - KV_H); - - fp_accumulate_with_fp_weight(out_tile, attn_weight_tile, w_tile); - } - } - - partial_sums[worker_id] = out_tile; - - memoryBarrierShared(); - barrier(); - - // Tree reduction to compute the overall result. - for (int i = NUM_WORKERS_PER_OUT / 2; i > 0; i /= 2) { - if (worker_id < i) { - accumulate_out_tile_with_out_tile( - partial_sums[worker_id], partial_sums[worker_id + i]); - } - memoryBarrierShared(); - barrier(); - } - - // Only the first thread will write out the result - if (worker_id == 0) { - out_tile = partial_sums[0]; - store_sdpa_out_tile_with_checks( - out_tile, - d4, - s, - q_h, - D4, - S, - Q_H); - } -} diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coop.yaml b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coop.yaml deleted file mode 100644 index 33ec2f8b322..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coop.yaml +++ /dev/null @@ -1,25 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -sdpa_compute_out_coop: - parameter_names_with_default_values: - DTYPE: float - IO_STORAGE: texture3d - V_CACHE_STORAGE: texture3d - TILE_K4: 1 - TILE_N4: 1 - generate_variant_forall: - combination: - parameter_names: [IO_STORAGE, V_CACHE_STORAGE] - combos: - - parameter_values: [texture3d, texture3d] - - parameter_values: [buffer, texture3d] - - parameter_values: [buffer, buffer] - DTYPE: - - VALUE: float - - VALUE: half - shader_variants: - - NAME: sdpa_compute_out_coop diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp index 4f6e79a2db3..001a94b236e 100644 --- a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp @@ -96,80 +96,35 @@ constexpr CoopmatTileDims kDq8caQ4gswCoopmatDims = {64, 32, 32, 128, 2}; // kernel rather than the actual mistake. Splitting them lets a wrong-family // token be rejected here, by name. // -// dbuf1-4 and the bare "tsweep_t" namespace exist for both families; the -// -tr/-zp/-zpn/-zpb variants are dq8ca-only (there is no q4gsw shader for -// any of them). +// 2026-09-15 housekeeping: removed the dead "dbuf1"/"dbuf2"/"dbuf3" and bare +// "tsweep_t" prefixes -- no GLSL for any of them ever existed in this branch +// (q4gsw or dq8ca), so they could only ever have produced a confusing +// shader-lookup crash instead of a clean validation error. "tsweep_dbuf4_t" +// is the only real q4gsw entry (matches the one shipped shader file). static const char* const kQ4gswTsweepPrefixes[] = { - "tsweep_dbuf1_t", - "tsweep_dbuf2_t", - "tsweep_dbuf3_t", "tsweep_dbuf4_t", - // (coopmat-lds-fence 2026-09-03: tsweep_dbuf4nf_t / tsweep_dbuf4zpgtrnf_t - // -- unfenced clones of the two shipped defaults, used to measure the cost - // of the memoryBarrierShared() fix via an interleaved same-binary A/B. - // Result: 8da4w +0.001%, 4w -0.002%, both inside a 0.017-0.037% noise band. - // Deleted after measurement.) - "tsweep_t", }; +// 2026-09-15 housekeeping: this branch now ships only the validated default +// (tsweep_dbuf4zpgtr_t). Removed alongside their now-deleted GLSL/YAML: +// "dbuf1"/"dbuf2"/"dbuf3" and bare "tsweep_t" (dead -- no GLSL ever existed +// for any of them, same as the q4gsw list above); "dbuf4" (the original +// 2026-08-26 default, twice superseded); "dbuf4zpg" (the 2026-08-28 default, +// superseded 2026-09-01 -- was also this function's invalid-token fallback, +// now repointed at the current default below); "dbuf4zpgtr3" (correctness- +// clean but +4.71% slower); "dbuf4zpgbw2"/"dbuf4zpgbw3" (correctness-clean, +// ~0% delta, no reason to keep as a build target). Full history (promotion +// dates, measured deltas, why each was rejected) preserved in the comments +// below and in the archived files themselves, on branch +// yanwen/release14-quant-shaders-archived-2026-09-15. (dbuf4zpgbv4/bw/bwr +// were already dead -- KNOWN-INCORRECT and deliberately never listed here -- +// their GLSL/YAML is removed too, same archive branch.) static const char* const kDq8caTsweepPrefixes[] = { - "tsweep_dbuf1_t", - "tsweep_dbuf2_t", - "tsweep_dbuf3_t", - "tsweep_dbuf4_t", - // zpi + compile-time elision of the statically-true a_active guard - // (intervention G of dq8ca-prefill-stall-reduction), combined with the - // dbuf4 default's own B_STRIDE_U32 skew removal + coalesced B-store - // rewrite. Superseded 2026-09-01 by tsweep_dbuf4zpgtr_t below -- kept - // listed (env-var-selectable) for comparison/rollback. - "tsweep_dbuf4zpg_t", // dbuf4zpg with its per-thread scalar A-staging replaced by a // coopMat-mediated coopMatLoad(global)->coopMatStore(LDS) sequence (B // staging/zp-hoist/nibble-widening unchanged). PROMOTED 2026-09-01 as // the shipped default -- see dq8ca_coopmat_variant() below. "tsweep_dbuf4zpgtr_t", - // dbuf4zpgtr with tr3's shared-A layout (scalar int8_t, row-major over - // the full chunk, A_ROW_PAD_I8 per-row pad). Correctness-clean 12/12 but - // MEASURED SLOWER (+4.71% best case) -- kept opt-in as a recorded - // negative result. Never a default. - "tsweep_dbuf4zpgtr3_t", - // B-staging ownership moved to one thread per 2-uint pair, all 256 threads - // active. bw2 keeps the uint array, bw3 retypes it to uvec2 for a real wide - // store. Both CORRECT (14/14 x3); neither faster (bw2 ~0%, bw3 +0.34%) -- - // the ISA shows scalar ds_store_b32 11->3 for nothing, because our LDS gap - // vs gemm-ubm TR3 is LOADS (80 vs 24), not stores (24 vs 8). OPT-IN only. - // (dbuf4zpgbv4/bw/bwr are KNOWN-INCORRECT -- deliberately NOT listed here - // so they cannot be selected; see their yaml headers for the isolated - // cause.) - "tsweep_dbuf4zpgbw2_t", - "tsweep_dbuf4zpgbw3_t", - // (dq8ca-dequant-unpack-ablation Addendum 11 -- abl_aconst/abl_areadc/ - // abl_abconst -- were measurement-only variants deleted once each - // attribution was recorded; see openspec/changes/dq8ca-dequant-unpack- - // ablation/results/README.md.) - // (dq8ca-dequant-unpack-ablation and its 2026-08-26 follow-ups on - // xgpusw-debug08 -- abl_nodq/abl_nonib/abl_both/abl_nolds/abl_bconst/ - // abl_bcont/abl_breadc/str4/str6/str8/bcoal -- were measurement-only - // variants deleted once each attribution was recorded; see - // openspec/changes/dq8ca-dequant-unpack-ablation/results/. Two real - // findings from that investigation WERE promoted to the shipped default: - // see the B_STRIDE_U32 comment (the LDS skew removal) and the - // BCoalIndex/bcoal_index comment (the coalesced B-store rewrite) in - // linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4.glsl.) - // 2026-08-28: removed the now-dead prefixes for "-tr"/"-trm"/"-trd" - // (row-major-A coopMatLoad staging, deprioritized -- see - // dq8ca-uvec4-coopmat-redesign/results.md: the reference kernel's read is - // narrow too, so this direction didn't hold), "-zp"/"-zpn"/"-zpb"/"-zpx" - // (superseded single-intervention isolations from dq8ca-prefill-stall- - // reduction, folded into "-zpg" above), and "-zpi"/"-zpk" (isolation - // variants used only to attribute G vs F in that investigation). Their - // shader files were deleted with them -- this branch (release14-quant- - // shaders) ships only the validated default; the experimental siblings - // and isolation variants live on dq8ca-uvec4-redesign/dq8ca-arch-redesign - // instead. Leaving a dead prefix here means an env var could reference a - // shader that no longer exists and crash confusingly at shader lookup - // instead of failing the validation check cleanly. - "tsweep_t", }; // (The measurement-only ablation variants and their prefix list lived here @@ -354,7 +309,11 @@ static const std::string& dq8ca_coopmat_variant() { if (is_dq8ca_shippable_token(v)) { return v; } - return std::string("tsweep_dbuf4zpg_t128x64k32g42s32"); + // 2026-09-15: was "tsweep_dbuf4zpg_t128x64k32g42s32" (the prior default); + // that shader was removed in the same housekeeping pass that pruned this + // fallback's dead alternatives above, so an invalid token now falls back + // to the current default instead. + return std::string("tsweep_dbuf4zpgtr_t128x64k32g42s32"); }(); return variant; } From 1731d590bae2468d5c98f53742e08bc7fb4c124b Mon Sep 17 00:00:00 2001 From: Yanwen Xu Date: Tue, 15 Sep 2026 12:28:18 -0700 Subject: [PATCH 27/28] [ET-VK] Housekeeping pass 2: remove dead qw_coopmat shaders, trim q4gsw tile-sweep yaml Second pass, same goal as the previous commit: this branch ships only the production shaders, minimal enough to serve as a future release branch. - linear_qw_coopmat.{glsl,yaml}, linear_dq8ca_qw_coopmat.{glsl,yaml}: removed. Zero real dispatch references anywhere in the codebase -- the only mentions were two comment lines in QuantizedLinear.cpp citing their tile geometry as historical provenance for two constants. Comments updated to explain what those constants are now that the files they cited are gone, and kDq8caQ4gswCoopmatDims's value corrected to the real current shipped tile (it held a stale, pre-08-26 geometry -- harmless in practice since it's an unreachable defensive fallback, but confusing to leave wrong). - linear_q4gsw_coopmat_tsweep_dbuf4.yaml: trimmed from 97 shader_variants (specs/041's full tile-sweep search space, ~48 geometries x storage combos) down to 3 -- the shipped tile's (t128x128k16g22s32) buffer/texture3d/ weight-buffer storage combos. Only this one tile is ever dispatched; the other 94 were pure sweep residue never reachable outside an explicit ET_VK_Q4GSW_COOPMAT_VARIANT override. Deliberately left alone: coopmat_mm.glsl / GemmCoopmat.h (Linear.cpp / Matmul.cpp's general fp16 coopmat path) -- real shared backend infrastructure, not one of our experimental variants, just permanently gated off on M51 via !is_integrated_gpu(). Out of scope for this cleanup. Rebuilt clean (install + custom_ops/test_llama_microbench) after this change; device re-verification not run for this commit. Full pre-cleanup history for everything removed today is preserved on branch yanwen/release14-quant-shaders-archived-2026-09-15 (snapshotted before the first housekeeping commit, so it also covers this one). --- .../ops/glsl/linear_dq8ca_qw_coopmat.glsl | 574 ------------------ .../ops/glsl/linear_dq8ca_qw_coopmat.yaml | 42 -- .../graph/ops/glsl/linear_qw_coopmat.glsl | 520 ---------------- .../graph/ops/glsl/linear_qw_coopmat.yaml | 47 -- 4 files changed, 1183 deletions(-) delete mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_qw_coopmat.glsl delete mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_qw_coopmat.yaml delete mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_qw_coopmat.glsl delete mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_qw_coopmat.yaml diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_qw_coopmat.glsl b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_qw_coopmat.glsl deleted file mode 100644 index 5edb05001c7..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_qw_coopmat.glsl +++ /dev/null @@ -1,574 +0,0 @@ -/* - * Copyright (c) Meta Platforms, Inc. and affiliates. - * All rights reserved. - * - * This source code is licensed under the BSD-style license found in the - * LICENSE file in the root directory of this source tree. - */ - -/* - * specs/025/026/027: loop structure updated from "dbuf4" (store-first, - * single-buffered-until-prefetch) to "dbuf2" (store-first, prefetch-first - * prologue) per specs/025 User Story 1's re-confirmed loop-structure winner - * for this shader; tile geometry updated from the prior 128x64/K32/2x2/s64 - * to 64x32/K32/1x2/s64 per specs/027's e2e-ranked sweep winner - * (tsweep_t64x32k32g12s64) -- confirmed +9.32% real end-to-end prefill - * throughput on M5 EVT1 (Llama 3.1 8B, 2048-token prefill), not just - * isolated-kernel GFLOP/s. See specs/027-e2e-tile-sweep/results/sweep-report.md. - * - * KHR Cooperative Matrix variant of the dynamically-quantized-activation - * linear tiled shader (WEIGHT_NBITS=4): - * 4 -> linear_dq8ca_q4gsw_coopmat INT4 group-symmetric weight - * - * Performs: out[M,N] = dequant(int8_act) * dequant(int_w) (+ bias) - * via coopmat x coopmat -> coopmat on the matrix unit. - * - * Math (per group; per-channel INT8 is the num_groups == 1 special case - * where the single "group" spans all of K): - * accum_int32 = sum_k(int8_in_k * int_w_signed_k) // coopMatMulAdd - * adjusted = accum_int32 - input_zp[m] * wsum_signed[group, n] - * delta_fp = float(adjusted) * (input_scale[m] * weight_scale[group, n]) - * result_fp += delta_fp // across groups - * - * Because INT4 weights are sign-extended to int8 in the B-stage, the - * "8 * input_sum" term of the tiled correction (which compensates for - * unsigned int4 nibbles in dotPacked4x8) cancels out and is not needed. - * - * Loop structure ("dbuf2", specs/023-8da4w-int8-dbuf-sweep naming): prologue - * prefetches chunk 0 into temp registers only (no shared-memory write, no - * barrier); each loop iteration then does store(temp -> cur slice) - * -> barrier() [UNCONDITIONAL, every iteration] -> MMA(cur) -> prefetch(next - * chunk -> temp) [skipped on the last chunk]. Iteration `chunk` stores the - * data FOR ITSELF (already prefetched by the previous iteration, or by the - * prologue for chunk 0), immediately before using it. The same inversion - * applies to the group wsum/wsc ping-pong: this variant stores the CURRENT - * group's values (prefetched by the previous group's last chunk) at the head - * of the group's first chunk. Group 0's wsum/wsc are unaffected -- set up - * directly in the prologue. The nested groups x chunks loop and - * unconditional group epilog are kept exactly as before -- flattening them - * with a conditional coopmat epilog crashes the Xclipse PAL compiler at - * large spec-resolved trip counts (specs/023 finding). - * - * Per-(group, N) weight sums/scales live in a SECOND ping-pong pair indexed - * by group parity: the next group's values are prefetched into registers - * and stored to the other wsum/wsc slice during the iteration that crosses - * the group boundary, and the regular per-iteration barrier makes them - * visible before that group's epilog runs. Per-row activation zp/scale - * broadcasts are group-invariant and loaded once in the prologue. - * - * LDS layout for the MMA operands: K-slab split + ColumnMajor B + per-col - * skew padding: the int8 WMMA matB lane layout wants 4 K-contiguous bytes - * per lane, so a RowMajor B in LDS forces per-byte ds_load + v_perm repack - * chains. ColumnMajor with a +1-uint skew per column gives one ds_load_b32 - * per lane with a bank-conflict-free col stride. Each uint holds 4 packed - * int8. - * - * Tile hierarchy (yaml): MMA 16x16x16 int8, WG_TILE 64x32, WG_TILE_K = 32, - * 2 subgroups x 64 threads (1x2 grid) -- specs/027's e2e-ranked winner. - * SUBGROUP_SIZE stays 64: specs/026 found subgroup=32 is legal (no compiler - * crash) but sharply tile-shape-dependently INCORRECT, and this tile shape - * was not one of the two shapes specs/026 found fully-correct at subgroup=32 - * -- see specs/026-8da4w-subgroup32-sweep/results/ for the full picture - * before considering subgroup=32 at this or any other tile shape. - * - * Hard preconditions: - * M % WG_TILE_M == 0, N % WG_TILE_N == 0, K % WG_TILE_K == 0, - * INT4: group_size % WG_TILE_K == 0, - * device exposes coopmatx-> at 16x16x16. - */ - -#version 450 core - -#extension GL_KHR_cooperative_matrix : require -#extension GL_KHR_memory_scope_semantics : require -#extension GL_KHR_shader_subgroup_basic : enable -#extension GL_EXT_shader_explicit_arithmetic_types : require -#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require -#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require -#extension GL_EXT_control_flow_attributes : enable - -#define PRECISION ${PRECISION} - -$if WEIGHT_NBITS == 4: - #define WEIGHT_INT4 - -$if HAS_BIAS: - #define HAS_BIAS - -$if WEIGHT_STORAGE == "buffer": - #define WEIGHT_BUFFER - -layout(std430) buffer; - -#include "common.glslh" - -// Bindings — match add_linear_dqa_qw_node arg order: -// output(0), fp_input(1), packed_int8_input(2), int_input_sums(3 - unused), -// input_scales(4), input_zps(5), packed_weight(6), weight_sums(7), -// weight_scales(8), bias(9). -${layout_declare_tensor(B, "w", "t_output", "half", "buffer", is_scalar_array=True)} -${layout_declare_tensor(B, "r", "t_input", "half", "buffer", is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_packed_int8_input", "int", "buffer", is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_int8_input_sums", "int", "buffer", is_scalar_array=True)} -${layout_declare_tensor(B, "r", "t_int8_input_scales", "half", "texture3d")} -${layout_declare_tensor(B, "r", "t_int8_input_zps", "int8", "texture3d")} -${layout_declare_tensor(B, "r", "t_packed_weight", "int", WEIGHT_STORAGE, is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_weight_sums", "int", "buffer", is_scalar_array=True)} -${layout_declare_tensor(B, "r", "t_weight_scales", "half", "buffer", is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_bias", "half", "buffer", is_scalar_array=True)} - -${layout_declare_ubo(B, "ivec4", "output_sizes")} -${layout_declare_ubo(B, "ivec4", "input_sizes")} - -layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; - -${layout_declare_spec_const(C, "int", "apply_bias", "0")} -// INT4 only; inert (0) for INT8 so the dispatcher's spec list lines up. -${layout_declare_spec_const(C, "int", "K4_per_group", "0")} -// Trip-count source for the coopmat K loop, passed as a spec constant (not -// derived from the runtime sizes UBO): the Xclipse/AMD-PAL shader compiler -// crashes (null deref in vkCreateComputePipelines) when a loop containing -// coopMatMulAdd has a UBO-derived trip count. INT4: number of quant groups; -// INT8: number of K-chunks. -// -// Unlike linear_qw_coopmat, this spec-const workaround is INTENTIONALLY kept -// here: on 2026-06-30 the UBO-direct method (sizes UBO feeding num_chunks/N -// directly) was A/B'd on this shader and produced wrong results for the -// coopmat (buffer) path at M>=128, while this spec-const version validated -// clean — see add_linear_dqa_qw_node in QuantizedLinear.cpp. -${layout_declare_spec_const(C, "int", "num_groups_arg", "0")} -// Output width N for coopMatStore: the Xclipse compiler MISCOMPILES -// coopMatStore whose offset/stride derive from a UBO value (only the first -// store per subgroup lands correctly; standalone repro cm_acc2). -${layout_declare_spec_const(C, "int", "out_N_arg", "0")} - -// Tile geometry -const uint MMA_M = ${MMA_M}; -const uint MMA_N = ${MMA_N}; -const uint MMA_K = ${MMA_K}; - -const uint WG_TILE_M = ${WG_TILE_M}; -const uint WG_TILE_N = ${WG_TILE_N}; -const uint WG_TILE_K = ${WG_TILE_K}; - -const uint SG_GRID_X = ${SG_GRID_X}; -const uint SG_GRID_Y = ${SG_GRID_Y}; -const uint SUBGROUP_SIZE = ${SUBGROUP_SIZE}; -const uint NUM_SUBGROUPS = SG_GRID_X * SG_GRID_Y; -const uint WG_SIZE = NUM_SUBGROUPS * SUBGROUP_SIZE; - -const uint SG_TILE_M = WG_TILE_M / SG_GRID_Y; -const uint SG_TILE_N = WG_TILE_N / SG_GRID_X; -const uint MMAS_PER_SG_M = SG_TILE_M / MMA_M; -const uint MMAS_PER_SG_N = SG_TILE_N / MMA_N; - -const uint A_SLAB_INT8 = WG_TILE_M * MMA_K; -const uint B_USEFUL_U32 = MMA_K / 4u; -const uint B_STRIDE_U32 = B_USEFUL_U32 + 1u; // +1 skew -const uint B_SLAB_U32 = WG_TILE_N * B_STRIDE_U32; -const uint NUM_K_SLABS = WG_TILE_K / MMA_K; - -const uint A_SLAB_U32 = A_SLAB_INT8 / 4u; -const uint A_STRIDE_U32 = MMA_K / 4u; - -// One ping-pong slice covers all K-slabs of one chunk. -const uint ASH_SLICE_U32 = NUM_K_SLABS * A_SLAB_U32; -const uint BSH_SLICE_U32 = NUM_K_SLABS * B_SLAB_U32; - -// Double-buffered MMA operand staging. -shared uint Ash_int8[2u * ASH_SLICE_U32]; -shared uint Bsh_int8[2u * BSH_SLICE_U32]; - -// Per-WG-tile-row activation params (loaded ONCE at WG start; constant -// across groups). -shared int izp_sh[WG_TILE_M]; // int32 (cast from int8 source) for broadcast -shared float ifs_sh[WG_TILE_M]; // float32 (cast from fp16 source) for broadcast - -// Per-(group, output-channel) weight params, ping-ponged by group parity. -// (For per-channel INT8 only slice 0 is ever used.) -shared int wsum_sh[2u * WG_TILE_N]; -shared float wsc_sh[2u * WG_TILE_N]; - -#ifdef HAS_BIAS -shared float bias_sh[WG_TILE_N]; -#endif - -// Running fp32 accumulator (across all groups). -coopmat - result[MMAS_PER_SG_M][MMAS_PER_SG_N]; - -// Per-group int32 MMA accumulator. -coopmat - accum_int32[MMAS_PER_SG_M][MMAS_PER_SG_N]; - -void main() { - const uvec2 tileID = uvec2(gl_WorkGroupID.xy); - const uvec2 warpInTile = uvec2( - gl_SubgroupID % SG_GRID_X, - gl_SubgroupID / SG_GRID_X); - - const uint K = uint(input_sizes.x); - const uint N = uint(output_sizes.x); - const uint N4 = (N + 3u) / 4u; - const uint nblocks_x_A = (K + 3u) >> 2u; - -#ifdef WEIGHT_INT4 - const uint num_groups = uint(num_groups_arg); - const uint CHUNKS_PER_GROUP = uint(K4_per_group) * 4u / WG_TILE_K; -#else - // Per-channel: a single quant "group" spanning all of K. The nested - // groups x chunks loop below collapses to a flat chunk loop, the wsum/wsc - // ping-pong never crosses a boundary, and the epilog runs exactly once. - const uint num_groups = 1u; - const uint CHUNKS_PER_GROUP = uint(num_groups_arg); -#endif - const uint num_chunks = num_groups * CHUNKS_PER_GROUP; - - const uint tile_m_start = WG_TILE_M * tileID.y; - const uint tile_n_start = WG_TILE_N * tileID.x; - - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - result[i][j] = coopmat(0.0); - accum_int32[i][j] = coopmat(0); - } - } - - // --- A staging thread map: one (m4, k4) ivec4 block per active thread --- - // (4 M-rows x 4 K-positions; each block expands to 4 slab-major LDS uints.) - const uint K_BLOCKS_PER_CHUNK = WG_TILE_K >> 2u; - const uint A_ACTIVE_THREADS = (WG_TILE_M >> 2u) * K_BLOCKS_PER_CHUNK; - const uint a_m_block = gl_LocalInvocationID.x / K_BLOCKS_PER_CHUNK; - const uint a_k_block = gl_LocalInvocationID.x % K_BLOCKS_PER_CHUNK; - const bool a_active = gl_LocalInvocationID.x < A_ACTIVE_THREADS; - -#ifdef WEIGHT_INT4 - // --- B staging thread map: (block, col) slots; each slot extracts one - // ColumnMajor LDS uint (4 K-contiguous sign-extended int8) --- - // INT4 weight block grid (see pack_q4_linear_weight.glsl): block (k4, n8) - // covers K=[k4*4, k4*4+3] x N=[n8*8, n8*8+7]. Within a block, int32[r] - // nibble col c maps to N = n8*8 + r + (c&1 ? 4 : 0), K = k4*4 + c/2 — one - // (component, parity) pair yields exactly the 4 K-contiguous bytes of one - // N column = one ColumnMajor LDS uint. - const uint B_TOTAL_SLOTS = K_BLOCKS_PER_CHUNK * WG_TILE_N; - const uint B_SLOTS_PER_THREAD = B_TOTAL_SLOTS / WG_SIZE; - const uint N8_PER_TILE = WG_TILE_N >> 3u; -#else - // --- B staging thread map: one (k4, n4) ivec4 block per active thread --- - // INT8 weight block layout: wblk[n_in_blk] packs 4 K-contiguous bytes for - // N-col (n4*4 + n_in_blk) — exactly one ColumnMajor LDS uint, written - // as-is (no byte repack). - const uint B_FETCH_SLOTS = K_BLOCKS_PER_CHUNK * (WG_TILE_N >> 2u); - const uint N4_PER_TILE = WG_TILE_N >> 2u; - const uint b_k4_in_chunk = gl_LocalInvocationID.x / N4_PER_TILE; - const uint b_n_uint_col = gl_LocalInvocationID.x % N4_PER_TILE; - const bool b_active = gl_LocalInvocationID.x < B_FETCH_SLOTS; -#endif - - // Prefetch temp registers. - ivec4 temp_A; -#ifdef WEIGHT_INT4 - ivec4 temp_B[B_SLOTS_PER_THREAD]; - int temp_wsum; - float temp_wsc; -#else - ivec4 temp_B; -#endif - - // ========================================================= - // PROLOGUE - // ========================================================= - // One-time: per-row input zp + scale (texture3d, one m4-block of 4 rows per - // texel) — constant across K groups. - if (gl_LocalInvocationID.x < (WG_TILE_M >> 2u)) { - const uint m4 = (tile_m_start >> 2u) + gl_LocalInvocationID.x; - const vec4 sc = vec4(texelFetch(t_int8_input_scales, ivec3(m4, 0, 0), 0)); - const ivec4 zp = texelFetch(t_int8_input_zps, ivec3(m4, 0, 0), 0); - const uint base = gl_LocalInvocationID.x * 4u; - ifs_sh[base + 0u] = sc.x; ifs_sh[base + 1u] = sc.y; - ifs_sh[base + 2u] = sc.z; ifs_sh[base + 3u] = sc.w; - izp_sh[base + 0u] = zp.x; izp_sh[base + 1u] = zp.y; - izp_sh[base + 2u] = zp.z; izp_sh[base + 3u] = zp.w; - } - // Group 0 weight sums/scales -> slice 0. - if (gl_LocalInvocationID.x < WG_TILE_N) { - const uint n_idx = tile_n_start + gl_LocalInvocationID.x; - f16vec4 sv = t_weight_scales[n_idx >> 2u]; - wsc_sh[gl_LocalInvocationID.x] = float(sv[n_idx & 3u]); - wsum_sh[gl_LocalInvocationID.x] = t_weight_sums[n_idx]; - } - memoryBarrierShared(); - barrier(); - - // izp/ifs are per-row activation params, constant across K groups — - // broadcast them into coopmats ONCE; the group epilog reuses them every - // group (they depend only on the row block i, not on the group or j). - coopmat - izp_bcast[MMAS_PER_SG_M]; - coopmat - ifs_bcast[MMAS_PER_SG_M]; - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - const uint local_m_base = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); - coopMatLoad( - izp_bcast[i], izp_sh, - local_m_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutColumnMajor); - coopMatLoad( - ifs_bcast[i], ifs_sh, - local_m_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutColumnMajor); - } - - // dbuf2: prefetch chunk 0 into temp registers only -- no shared-memory - // write, no barrier here. The main loop's first iteration stores temp - // into slice 0 and barriers as normal (uniform code path for every chunk, - // including chunk 0). - if (a_active) { - const uint m4_global = (tile_m_start >> 2u) + a_m_block; - temp_A = t_packed_int8_input[m4_global * nblocks_x_A + a_k_block]; - } -#ifdef WEIGHT_INT4 - [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { - const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; - const uint block_in_chunk = slot >> 3u; - const uint k4_blk = block_in_chunk / N8_PER_TILE; - const uint n8_blk = (tile_n_start >> 3u) + (block_in_chunk % N8_PER_TILE); -#ifdef WEIGHT_BUFFER - temp_B[si] = t_packed_weight[(n8_blk * nblocks_x_A) + k4_blk]; -#else - temp_B[si] = texelFetch(t_packed_weight, ivec2(k4_blk, n8_blk), 0); -#endif - } -#else - if (b_active) { - const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; -#ifdef WEIGHT_BUFFER - temp_B = t_packed_weight[(b_k4_in_chunk * N4) + block_x_w]; -#else - temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, b_k4_in_chunk), 0); -#endif - } -#endif - - // ========================================================= - // MAIN LOOP — nested groups x chunks (the flattened single loop with a - // conditional coopmat epilog crashes the Xclipse PAL compiler at large - // spec-resolved trip counts). One barrier per chunk, UNCONDITIONAL every - // iteration (dbuf2 "store-first" ordering). Chunk iteration (global index - // `chunk`): - // 1. store — temp (already holding this chunk's data, from the - // prologue for chunk 0 or from the previous iteration's - // step 4) -> A/B slice (chunk%2), unpacking the weight; - // on a group boundary (first chunk of a group > 0), - // also store this group's wsum/wsc -> slice (group_i%2). - // 2. barrier — A/B slice (chunk%2) (and, on a group boundary, wsum/wsc - // slice (group_i%2)) fully written; UNCONDITIONAL, not - // skipped on the last chunk. - // 3. int8 MMA — on slice (chunk%2) into accum_int32. - // 4. prefetch — chunk+1 (A blocks, B blocks) into temp; when chunk+1 - // starts a new group, also that group's wsum/wsc - // element. Skipped entirely on the final chunk. - // The group epilog runs unconditionally at the tail of each group. - // ========================================================= - uint chunk = 0; - for (uint group_i = 0; group_i < num_groups; ++group_i) { - for (uint inner = 0; inner < CHUNKS_PER_GROUP; ++inner, ++chunk) { - const bool has_next = chunk + 1u < num_chunks; - const uint cur_a = (chunk % 2u) * ASH_SLICE_U32; - const uint cur_b = (chunk % 2u) * BSH_SLICE_U32; - - // --- 1. store temp (this chunk) -> cur slice --- - if (a_active) { - const uint slab_idx = a_k_block / (MMA_K >> 2u); - const uint k_uint_in_slab = a_k_block % (MMA_K >> 2u); - const uint base_row = a_m_block * 4u; - [[unroll]] for (uint m4i = 0; m4i < 4u; ++m4i) { - Ash_int8[cur_a + slab_idx * A_SLAB_U32 + (base_row + m4i) * A_STRIDE_U32 + k_uint_in_slab] = - uint(temp_A[m4i]); - } - } -#ifdef WEIGHT_INT4 - [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { - const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; - const uint block_in_chunk = slot >> 3u; - const uint col_in_block = slot & 7u; - const uint k4_in_chunk = block_in_chunk / N8_PER_TILE; - const uint n8_in_tile = block_in_chunk % N8_PER_TILE; - const uint r = col_in_block & 3u; - const uint parity = col_in_block >> 2u; - const int w = temp_B[si][r]; - const int base = int(4u * parity); - const int v0 = (((w >> (base + 0)) & 0xF) - 8) & 0xFF; - const int v1 = (((w >> (base + 8)) & 0xF) - 8) & 0xFF; - const int v2 = (((w >> (base + 16)) & 0xF) - 8) & 0xFF; - const int v3 = (((w >> (base + 24)) & 0xF) - 8) & 0xFF; - const uint n_col = n8_in_tile * 8u + r + parity * 4u; - const uint slab_idx = k4_in_chunk / (MMA_K >> 2u); - const uint k4_in_slab = k4_in_chunk % (MMA_K >> 2u); - Bsh_int8[cur_b + slab_idx * B_SLAB_U32 + n_col * B_STRIDE_U32 + k4_in_slab] = - uint(v0 | (v1 << 8) | (v2 << 16) | (v3 << 24)); - } - // Group boundary (this is the first chunk of a group other than - // group 0): store this group's wsum/wsc, prefetched by the previous - // group's last chunk (step 4 below). - if (inner == 0u && group_i > 0u && gl_LocalInvocationID.x < WG_TILE_N) { - const uint wbase_cur = (group_i % 2u) * WG_TILE_N; - wsum_sh[wbase_cur + gl_LocalInvocationID.x] = temp_wsum; - wsc_sh[wbase_cur + gl_LocalInvocationID.x] = temp_wsc; - } -#else - if (b_active) { - const uint slab_idx = b_k4_in_chunk / (MMA_K >> 2u); - const uint k4_in_slab = b_k4_in_chunk % (MMA_K >> 2u); - const uint n_col_base = b_n_uint_col * 4u; - [[unroll]] for (uint n_in_blk = 0u; n_in_blk < 4u; ++n_in_blk) { - Bsh_int8[cur_b + slab_idx * B_SLAB_U32 + (n_col_base + n_in_blk) * B_STRIDE_U32 + k4_in_slab] = - uint(temp_B[n_in_blk]); - } - } -#endif - - // --- 2. barrier — cur slice(s) fully written --- - memoryBarrierShared(); - barrier(); - - // --- 3. int8 MMA on the cur slice --- - [[unroll]] for (uint k = 0; k < NUM_K_SLABS; ++k) { - const uint slab_a_base_u32 = cur_a + k * A_SLAB_U32; - const uint slab_b_base_u32 = cur_b + k * B_SLAB_U32; - - coopmat matA[MMAS_PER_SG_M]; - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - const uint row_a = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); - coopMatLoad( - matA[i], Ash_int8, - slab_a_base_u32 + row_a * A_STRIDE_U32, - A_STRIDE_U32, - gl_CooperativeMatrixLayoutRowMajor); - } - - coopmat matB; - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint col_b = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopMatLoad( - matB, Bsh_int8, - slab_b_base_u32 + col_b * B_STRIDE_U32, - B_STRIDE_U32, - gl_CooperativeMatrixLayoutColumnMajor); - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - accum_int32[i][j] = coopMatMulAdd(matA[i], matB, accum_int32[i][j]); - } - } - } - - // --- 4. prefetch chunk+1 -> temp --- - if (has_next) { - const bool group_crossing = (inner + 1u == CHUNKS_PER_GROUP); - const uint chunkK_nxt = (chunk + 1u) * WG_TILE_K; - if (a_active) { - const uint m4_global = (tile_m_start >> 2u) + a_m_block; - const uint k4_global = (chunkK_nxt >> 2u) + a_k_block; - temp_A = t_packed_int8_input[m4_global * nblocks_x_A + k4_global]; - } -#ifdef WEIGHT_INT4 - [[unroll]] for (uint si = 0; si < B_SLOTS_PER_THREAD; ++si) { - const uint slot = gl_LocalInvocationID.x + si * WG_SIZE; - const uint block_in_chunk = slot >> 3u; - const uint k4_blk = (chunkK_nxt >> 2u) + block_in_chunk / N8_PER_TILE; - const uint n8_blk = (tile_n_start >> 3u) + (block_in_chunk % N8_PER_TILE); -#ifdef WEIGHT_BUFFER - temp_B[si] = t_packed_weight[(n8_blk * nblocks_x_A) + k4_blk]; -#else - temp_B[si] = texelFetch(t_packed_weight, ivec2(k4_blk, n8_blk), 0); -#endif - } - if (group_crossing && gl_LocalInvocationID.x < WG_TILE_N) { - const uint n_idx = tile_n_start + gl_LocalInvocationID.x; - f16vec4 sv = t_weight_scales[(group_i + 1u) * N4 + (n_idx >> 2u)]; - temp_wsc = float(sv[n_idx & 3u]); - temp_wsum = t_weight_sums[(group_i + 1u) * N + n_idx]; - } -#else - if (b_active) { - const uint block_y_w = (chunkK_nxt >> 2u) + b_k4_in_chunk; - const uint block_x_w = (tile_n_start >> 2u) + b_n_uint_col; -#ifdef WEIGHT_BUFFER - temp_B = t_packed_weight[(block_y_w * N4) + block_x_w]; -#else - temp_B = texelFetch(t_packed_weight, ivec2(block_x_w, block_y_w), 0); -#endif - } -#endif - } - } // chunks - - // --- Group epilog: dequant accum_int32 -> result, reset accum --- - { - const uint wbase = (group_i % 2u) * WG_TILE_N; - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint local_n_base = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - - coopmat wsum_bcast; - coopMatLoad( - wsum_bcast, wsum_sh, - wbase + local_n_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutRowMajor); - - coopmat wsc_bcast; - coopMatLoad( - wsc_bcast, wsc_sh, - wbase + local_n_base, /*stride=*/0u, - gl_CooperativeMatrixLayoutRowMajor); - - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - coopmat adjusted = - accum_int32[i][j] - izp_bcast[i] * wsum_bcast; - coopmat adjusted_fp = - coopmat(adjusted); - coopmat scales_outer = - ifs_bcast[i] * wsc_bcast; - result[i][j] += adjusted_fp * scales_outer; - accum_int32[i][j] = coopmat(0); - } - } - } - } // groups - - // --- Bias (optional) --- -#ifdef HAS_BIAS - if (apply_bias > 0) { - for (uint t = gl_LocalInvocationID.x; t < WG_TILE_N; t += WG_SIZE) { - bias_sh[t] = float(t_bias[tile_n_start + t]); - } - memoryBarrierShared(); - barrier(); - } -#endif - - // --- Store result tile --- - // N for the store address math MUST come from the spec constant, not the - // sizes UBO (see out_N_arg above). - const uint N_out = uint(out_N_arg); - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint gi = tile_m_start + MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); - const uint gj = tile_n_start + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - -#ifdef HAS_BIAS - if (apply_bias > 0) { - const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopmat bias_tile; - coopMatLoad(bias_tile, bias_sh, local_n, 0u, gl_CooperativeMatrixLayoutRowMajor); - result[i][j] += bias_tile; - } -#endif - - coopmat out_tile = - coopmat(result[i][j]); - coopMatStore( - out_tile, t_output, - gi * N_out + gj, N_out, - gl_CooperativeMatrixLayoutRowMajor); - } - } -} diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_qw_coopmat.yaml b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_qw_coopmat.yaml deleted file mode 100644 index 7bc39d224a8..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_qw_coopmat.yaml +++ /dev/null @@ -1,42 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -# coopmat x coopmat -> coopmat variant of the -# dynamically-quantized-activation linear tiled shader (INT4 group-symmetric -# weight). -# WEIGHT_NBITS=4 -> linear_dq8ca_q4gsw_coopmat (INT4 group-symmetric) -# Requires the VK_COMPONENT_TYPE_SINT8_KHR cooperative matrix property to be -# enumerated on the device. -# specs/027-e2e-tile-sweep: tile geometry updated from 128x64/K32/2x2 to -# 64x32/K32/1x2 (specs/027's e2e-ranked sweep winner, +9.32% confirmed real -# end-to-end prefill throughput vs the prior geometry, not just isolated -# microbenchmark GFLOP/s). SUBGROUP_SIZE stays 64: specs/026 found -# subgroup=32 is legal (no compiler crash) but sharply tile-shape-dependently -# incorrect, and this tile shape was not verified correct at subgroup=32 -- -# see specs/026-8da4w-subgroup32-sweep/results/ before considering it. - -linear_dq8ca_qw_coopmat: - parameter_names_with_default_values: - PRECISION: highp - HAS_BIAS: false - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - MMA_M: 16 - MMA_N: 16 - MMA_K: 16 - WG_TILE_M: 64 - WG_TILE_N: 32 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - shader_variants: - - NAME: linear_dq8ca_q4gsw_coopmat_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - - NAME: linear_dq8ca_q4gsw_coopmat_buffer_buffer_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: buffer diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_qw_coopmat.glsl b/backends/vulkan/runtime/graph/ops/glsl/linear_qw_coopmat.glsl deleted file mode 100644 index 1f9707c3b79..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_qw_coopmat.glsl +++ /dev/null @@ -1,520 +0,0 @@ -/* - * Copyright (c) Meta Platforms, Inc. and affiliates. - * All rights reserved. - * - * This source code is licensed under the BSD-style license found in the - * LICENSE file in the root directory of this source tree. - */ - -/* - * KHR Cooperative Matrix variant of the weight-only int4 quantized linear - * tiled shader (WEIGHT_NBITS=4 in the yaml): - * 4 -> linear_q4gsw_coopmat INT4 group-symmetric weight - * (group_size = 4 * K4_per_group) - * - * Performs: out[M,N] = activation[M,K] * weight^T[N,K] (+ bias) - * - * Inner-loop math is pure fp16 -> fp32 MMA via coopMatMulAdd for both - * formats. The weight scale is applied during the B-tile store to shared - * memory: each int weight is unpacked (nibble - 8 for INT4; bitfieldExtract - * for INT8), cast to fp16, and multiplied by its scale before it lands in - * Bsh, keeping the K-loop a clean fp16 MMA. - * - * Loop structure follows the NVIDIA double-buffered GEMM reference - * (shmem_double_buf4.comp, "store-first" variant; see coopmat_mm_ref.glsl - * in test/custom_ops — measured 1.5x faster than the previous - * single-buffered skeleton at fp16 on Xclipse 970): - * - PROLOGUE: prefetch tile 0 from global memory into temp registers, then - * store it to shared-memory slice 0 (no barrier). - * - Each iteration: barrier -> global prefetch of the NEXT tile into temp - * -> MMA math on the CURRENT slice -> store temp into the OTHER slice. - * One barrier per iteration; the prefetch loads are in flight during the - * math and are only consumed at the store stage. - * - Ping-pong shared-memory slices make the overlap safe. - * - * Each thread keeps its 8 weight scales (2 f16vec4) in registers. For INT4 - * they are reloaded from global only when the prefetched chunk crosses a - * group boundary (a workgroup-uniform branch); for INT8 (per-channel = a - * single group spanning all of K) they are loaded once in the prologue. - * There is no scales staging in shared memory and no extra barrier. - * - * Tile hierarchy (yaml; mirrors the double-buffered reference): - * MMA_* per-MMA-instruction shape (16x16x16 fp16) - * WG_TILE_* output tile per workgroup (128x128) - * SG_GRID_* subgroup grid inside workgroup (4x2 = 8 subgroups) - * SUBGROUP_SIZE 32, forced at pipeline creation via the - * REQUIRED_SUBGROUP_SIZE annotation below - * - * Storage: activation/output forced to buffer; INT weight = texture2d or - * buffer (yaml variant). DTYPE = half only. - * - * Hard preconditions (no shape/alignment checks inside the shader): - * M % WG_TILE_M == 0 - * N % WG_TILE_N == 0 - * K % WG_TILE_K == 0 - * INT4: group_size % WG_TILE_K == 0 (each group = whole number of chunks) - * Misaligned shapes silently miscompute / overrun — gate at dispatch time. - */ - -// REQUIRED_SUBGROUP_SIZE = 32 - -#version 450 core - -#extension GL_KHR_cooperative_matrix : require -#extension GL_KHR_memory_scope_semantics : require -#extension GL_KHR_shader_subgroup_basic : enable -#extension GL_EXT_shader_explicit_arithmetic_types : require -#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require -#extension GL_EXT_control_flow_attributes : enable - -#define PRECISION ${PRECISION} - -$if WEIGHT_NBITS == 4: - #define WEIGHT_INT4 - -$if HAS_BIAS: - #define HAS_BIAS - -$if WEIGHT_STORAGE == "buffer": - #define WEIGHT_BUFFER - -layout(std430) buffer; - -#include "common.glslh" - -// Bindings — match the order used by add_linear_qw_node so the dispatch -// site can reuse the same arg layout. -${layout_declare_tensor(B, "w", "t_output", "half", "buffer", is_scalar_array=True)} -${layout_declare_tensor(B, "r", "t_input", "half", "buffer", is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_packed_weight", "int", WEIGHT_STORAGE, is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_weight_scales", "half", "buffer", is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_bias", "half", "buffer", is_scalar_array=True)} - -${layout_declare_ubo(B, "ivec4", "output_sizes")} -${layout_declare_ubo(B, "ivec4", "input_sizes")} - -layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; - -${layout_declare_spec_const(C, "int", "apply_bias", "0")} -// INT4 only; inert (0) for INT8 so the dispatcher's spec list lines up. -${layout_declare_spec_const(C, "int", "K4_per_group", "0")} -// Trip-count source for the coopmat K loop, passed as a spec constant (not -// derived from the runtime sizes UBO): the Xclipse/AMD-PAL shader compiler -// crashes (null deref in vkCreateComputePipelines) when a loop containing -// coopMatMulAdd has a UBO-derived trip count. INT4: number of quant groups; -// INT8: number of K-chunks. -${layout_declare_spec_const(C, "int", "num_groups_arg", "0")} -// Output width N for coopMatStore, as a spec constant: the same compiler -// MISCOMPILES coopMatStore whose offset/stride derive from a UBO value (only -// the first store per subgroup lands correctly; standalone repro cm_acc2). -${layout_declare_spec_const(C, "int", "out_N_arg", "0")} - -// --- Tile geometry (from yaml; defaults match coopmat_mm_ref) --- -const uint MMA_M = ${MMA_M}; -const uint MMA_N = ${MMA_N}; -const uint MMA_K = ${MMA_K}; - -const uint WG_TILE_M = ${WG_TILE_M}; -const uint WG_TILE_N = ${WG_TILE_N}; -const uint WG_TILE_K = ${WG_TILE_K}; - -const uint SG_GRID_X = ${SG_GRID_X}; -const uint SG_GRID_Y = ${SG_GRID_Y}; -const uint SUBGROUP_SIZE = ${SUBGROUP_SIZE}; -const uint NUM_SUBGROUPS = SG_GRID_X * SG_GRID_Y; -const uint WG_SIZE = NUM_SUBGROUPS * SUBGROUP_SIZE; - -const uint SG_TILE_M = WG_TILE_M / SG_GRID_Y; -const uint SG_TILE_N = WG_TILE_N / SG_GRID_X; -const uint MMAS_PER_SG_M = SG_TILE_M / MMA_M; -const uint MMAS_PER_SG_N = SG_TILE_N / MMA_N; - -// fp16: 8 elements per uvec4 (128-bit) -const uint FP16_PER_VEC4 = 8; -const uint A_STRIDE_VEC4 = (WG_TILE_K + FP16_PER_VEC4) / FP16_PER_VEC4; -const uint B_STRIDE_VEC4 = (WG_TILE_N + FP16_PER_VEC4) / FP16_PER_VEC4; - -// One ping-pong slice of each shared-memory buffer (in uvec4 units). -const uint ASH_SLICE = WG_TILE_M * A_STRIDE_VEC4; -const uint BSH_SLICE = WG_TILE_K * B_STRIDE_VEC4; - -// Double-buffered shared memory. -shared uvec4 Ash[2 * ASH_SLICE]; -shared uvec4 Bsh[2 * BSH_SLICE]; -#ifdef HAS_BIAS -shared float bias_sh[WG_TILE_N]; -#endif - -// Staging thread maps: each thread covers one uvec4 (8 fp16) per pass. -const uint INVS_PER_ROW_A = WG_TILE_K / FP16_PER_VEC4; -const uint A_ROWS_PER_PASS = WG_SIZE / INVS_PER_ROW_A; -const uint A_PASSES = WG_TILE_M / A_ROWS_PER_PASS; -const uint INVS_PER_ROW_B = WG_TILE_N / FP16_PER_VEC4; -const uint B_ROWS_PER_PASS = WG_SIZE / INVS_PER_ROW_B; -const uint B_PASSES = WG_TILE_K / B_ROWS_PER_PASS; - -// Fp32 accumulator coopmats (MMAS_PER_SG_M x MMAS_PER_SG_N per thread) -coopmat - result[MMAS_PER_SG_M][MMAS_PER_SG_N]; - -#ifdef WEIGHT_INT4 - -// Dequant one packed INT4 block column-pair into 8 scaled fp16 weights -// (one Bsh uvec4). col_lo/col_hi select the K row within the block. -uvec4 dequant_block( - const ivec4 wb, - const uint col_lo, - const uint col_hi, - const f16vec4 s0, - const f16vec4 s1) { - f16vec4 v0; - v0.x = float16_t(int(((wb[0] >> (4 * col_lo)) & 0xF)) - 8) * s0.x; - v0.y = float16_t(int(((wb[1] >> (4 * col_lo)) & 0xF)) - 8) * s0.y; - v0.z = float16_t(int(((wb[2] >> (4 * col_lo)) & 0xF)) - 8) * s0.z; - v0.w = float16_t(int(((wb[3] >> (4 * col_lo)) & 0xF)) - 8) * s0.w; - f16vec4 v1; - v1.x = float16_t(int(((wb[0] >> (4 * col_hi)) & 0xF)) - 8) * s1.x; - v1.y = float16_t(int(((wb[1] >> (4 * col_hi)) & 0xF)) - 8) * s1.y; - v1.z = float16_t(int(((wb[2] >> (4 * col_hi)) & 0xF)) - 8) * s1.z; - v1.w = float16_t(int(((wb[3] >> (4 * col_hi)) & 0xF)) - 8) * s1.w; - return uvec4( - packFloat2x16(v0.xy), packFloat2x16(v0.zw), - packFloat2x16(v1.xy), packFloat2x16(v1.zw)); -} - -#else // INT8 - -// Dequant 8 int8 weights (two ivec4 blocks, one K-row selected by shift) -// into 8 scaled fp16 weights (one Bsh uvec4). -uvec4 dequant_block( - const ivec4 wa, - const ivec4 wb, - const int shift, - const f16vec4 s0, - const f16vec4 s1) { - f16vec4 v0; - v0.x = float16_t(bitfieldExtract(wa.x, shift, 8)) * s0.x; - v0.y = float16_t(bitfieldExtract(wa.y, shift, 8)) * s0.y; - v0.z = float16_t(bitfieldExtract(wa.z, shift, 8)) * s0.z; - v0.w = float16_t(bitfieldExtract(wa.w, shift, 8)) * s0.w; - f16vec4 v1; - v1.x = float16_t(bitfieldExtract(wb.x, shift, 8)) * s1.x; - v1.y = float16_t(bitfieldExtract(wb.y, shift, 8)) * s1.y; - v1.z = float16_t(bitfieldExtract(wb.z, shift, 8)) * s1.z; - v1.w = float16_t(bitfieldExtract(wb.w, shift, 8)) * s1.w; - return uvec4( - packFloat2x16(v0.xy), packFloat2x16(v0.zw), - packFloat2x16(v1.xy), packFloat2x16(v1.zw)); -} - -#endif // WEIGHT_INT4 - -void main() { - const uvec2 tileID = uvec2(gl_WorkGroupID.xy); - const uvec2 warpInTile = uvec2( - gl_SubgroupID % SG_GRID_X, - gl_SubgroupID / SG_GRID_X); - - const uint K = uint(input_sizes.x); - const uint K4 = (K + 3u) / 4u; - const uint N4 = (uint(output_sizes.x) + 3u) / 4u; - -#ifdef WEIGHT_INT4 - const uint CHUNKS_PER_GROUP = uint(K4_per_group) * 4u / WG_TILE_K; - const uint num_chunks = uint(num_groups_arg) * CHUNKS_PER_GROUP; -#else - const uint num_chunks = uint(num_groups_arg); -#endif - - const uint tile_m_start = WG_TILE_M * tileID.y; - const uint tile_n_start = WG_TILE_N * tileID.x; - - // Initialize fp32 accumulators to zero. - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - result[i][j] = coopmat(0.0); - } - } - - const uint a_col = gl_LocalInvocationID.x % INVS_PER_ROW_A; - const uint a_row_offset = gl_LocalInvocationID.x / INVS_PER_ROW_A; - const uint b_col = gl_LocalInvocationID.x % INVS_PER_ROW_B; - const uint b_row_offset = gl_LocalInvocationID.x / INVS_PER_ROW_B; - -#ifdef WEIGHT_INT4 - // INT4 weight block grid (see pack_q4_linear_weight.glsl): block (k4, n8) - // covers K=[k4*4, k4*4+3] x N=[n8*8, n8*8+7]; buffer pitch = K4 blocks per - // n8 row, texture coord = ivec2(x=k4, y=n8). This thread's 8 N-values at - // any K-row live in column n8_blk of the block grid: - const uint n8_blk = (tile_n_start + b_col * 8u) >> 3u; - - // The K row within a block depends only on (b_row_offset & 3): chunkK and - // the pass offset are both multiples of 4. - const uint col_lo = 2u * (b_row_offset & 3u); - const uint col_hi = col_lo + 1u; - - // Per-thread per-group weight scales (8 consecutive N), kept in registers - // and reloaded only when the prefetched chunk crosses a group boundary. - const uint sc_n4 = (tile_n_start + b_col * 8u) >> 2u; - uint cached_group = 0xFFFFFFFFu; - f16vec4 sc0; - f16vec4 sc1; - - // Temp registers holding the prefetched (next) tile. - uvec4 temp_A[A_PASSES]; - ivec4 temp_B[B_PASSES]; // raw packed INT4 blocks; dequant at the store stage -#else - // INT8 weight block layout: t_packed_weight[k4 * N4 + n4] = ivec4 whose - // component n_in_blk packs 4 K-bytes (K of block k4) for N-col - // (n4*4 + n_in_blk). This thread's 8 N-values span two adjacent n4 blocks: - const uint n4_a = (tile_n_start + b_col * 8u) >> 2u; // n_start mult of 8 -> even - - // The byte within a packed uint depends only on (b_row_offset & 3): chunkK - // and the pass offset are both multiples of 4. - const int b_shift = int(8u * (b_row_offset & 3u)); - - // Per-thread per-channel weight scales (8 consecutive N), cached ONCE. - f16vec4 sc0 = t_weight_scales[n4_a]; - f16vec4 sc1 = t_weight_scales[n4_a + 1u]; - - // Temp registers holding the prefetched (next) tile. - uvec4 temp_A[A_PASSES]; - ivec4 temp_Ba[B_PASSES]; // raw packed INT8 blocks; dequant at the store stage - ivec4 temp_Bb[B_PASSES]; -#endif - - // ========================================================= - // PROLOGUE: prefetch chunk 0 into temp registers, then store to slice 0. - // ========================================================= - { - [[unroll]] for (uint p = 0; p < A_PASSES; ++p) { - const uint row = tile_m_start + p * A_ROWS_PER_PASS + a_row_offset; - const uint k_hv4 = (a_col * FP16_PER_VEC4) / 4u; - f16vec4 v0 = t_input[row * K4 + k_hv4]; - f16vec4 v1 = t_input[row * K4 + k_hv4 + 1u]; - temp_A[p] = uvec4( - packFloat2x16(v0.xy), packFloat2x16(v0.zw), - packFloat2x16(v1.xy), packFloat2x16(v1.zw)); - } -#ifdef WEIGHT_INT4 - [[unroll]] for (uint p = 0; p < B_PASSES; ++p) { - const uint k_row = p * B_ROWS_PER_PASS + b_row_offset; -#ifdef WEIGHT_BUFFER - temp_B[p] = t_packed_weight[n8_blk * K4 + (k_row >> 2u)]; -#else - temp_B[p] = texelFetch(t_packed_weight, ivec2(k_row >> 2u, n8_blk), 0); -#endif - } - cached_group = 0u; - sc0 = t_weight_scales[sc_n4]; - sc1 = t_weight_scales[sc_n4 + 1u]; -#else - [[unroll]] for (uint p = 0; p < B_PASSES; ++p) { - const uint k4 = (p * B_ROWS_PER_PASS + b_row_offset) >> 2u; -#ifdef WEIGHT_BUFFER - temp_Ba[p] = t_packed_weight[k4 * N4 + n4_a]; - temp_Bb[p] = t_packed_weight[k4 * N4 + n4_a + 1u]; -#else - temp_Ba[p] = texelFetch(t_packed_weight, ivec2(n4_a, k4), 0); - temp_Bb[p] = texelFetch(t_packed_weight, ivec2(n4_a + 1u, k4), 0); -#endif - } -#endif - } - { - [[unroll]] for (uint p = 0; p < A_PASSES; ++p) { - Ash[(p * A_ROWS_PER_PASS + a_row_offset) * A_STRIDE_VEC4 + a_col] = temp_A[p]; - } - [[unroll]] for (uint p = 0; p < B_PASSES; ++p) { -#ifdef WEIGHT_INT4 - Bsh[(p * B_ROWS_PER_PASS + b_row_offset) * B_STRIDE_VEC4 + b_col] = - dequant_block(temp_B[p], col_lo, col_hi, sc0, sc1); -#else - Bsh[(p * B_ROWS_PER_PASS + b_row_offset) * B_STRIDE_VEC4 + b_col] = - dequant_block(temp_Ba[p], temp_Bb[p], b_shift, sc0, sc1); -#endif - } - } - - // ========================================================= - // MAIN LOOP — one barrier per iteration. Iteration `chunk` does: - // 1. barrier — slice (chunk%2) fully written - // 2. prefetch — chunk+1 from global into temp (in flight during math) - // 3. MMA math — on slice (chunk%2) - // 4. store — temp (chunk+1, dequantized) into slice ((chunk+1)%2) - // ========================================================= - uint chunk; - for (chunk = 0; chunk + 1u < num_chunks; ++chunk) { - const uint cur_base_A = (chunk % 2u) * ASH_SLICE; - const uint cur_base_B = (chunk % 2u) * BSH_SLICE; - const uint nxt_base_A = ((chunk + 1u) % 2u) * ASH_SLICE; - const uint nxt_base_B = ((chunk + 1u) % 2u) * BSH_SLICE; - - barrier(); - - // --- prefetch chunk+1 -> temp --- - { - const uint chunkK_nxt = (chunk + 1u) * WG_TILE_K; - - [[unroll]] for (uint p = 0; p < A_PASSES; ++p) { - const uint row = tile_m_start + p * A_ROWS_PER_PASS + a_row_offset; - const uint k_hv4 = (chunkK_nxt + a_col * FP16_PER_VEC4) / 4u; - f16vec4 v0 = t_input[row * K4 + k_hv4]; - f16vec4 v1 = t_input[row * K4 + k_hv4 + 1u]; - temp_A[p] = uvec4( - packFloat2x16(v0.xy), packFloat2x16(v0.zw), - packFloat2x16(v1.xy), packFloat2x16(v1.zw)); - } -#ifdef WEIGHT_INT4 - [[unroll]] for (uint p = 0; p < B_PASSES; ++p) { - const uint k_row = chunkK_nxt + p * B_ROWS_PER_PASS + b_row_offset; -#ifdef WEIGHT_BUFFER - temp_B[p] = t_packed_weight[n8_blk * K4 + (k_row >> 2u)]; -#else - temp_B[p] = texelFetch(t_packed_weight, ivec2(k_row >> 2u, n8_blk), 0); -#endif - } - const uint group_nxt = (chunk + 1u) / CHUNKS_PER_GROUP; - if (group_nxt != cached_group) { - cached_group = group_nxt; - sc0 = t_weight_scales[group_nxt * N4 + sc_n4]; - sc1 = t_weight_scales[group_nxt * N4 + sc_n4 + 1u]; - } -#else - [[unroll]] for (uint p = 0; p < B_PASSES; ++p) { - const uint k4 = (chunkK_nxt + p * B_ROWS_PER_PASS + b_row_offset) >> 2u; -#ifdef WEIGHT_BUFFER - temp_Ba[p] = t_packed_weight[k4 * N4 + n4_a]; - temp_Bb[p] = t_packed_weight[k4 * N4 + n4_a + 1u]; -#else - temp_Ba[p] = texelFetch(t_packed_weight, ivec2(n4_a, k4), 0); - temp_Bb[p] = texelFetch(t_packed_weight, ivec2(n4_a + 1u, k4), 0); -#endif - } -#endif - } - - // --- MMA math on the cur slice --- - [[unroll]] for (uint k = 0; k < WG_TILE_K / MMA_K; ++k) { - const uint k_start = MMA_K * k; - - coopmat matA[MMAS_PER_SG_M]; - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - const uint row_a = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); - coopMatLoad( - matA[i], Ash, - cur_base_A + row_a * A_STRIDE_VEC4 + k_start / FP16_PER_VEC4, - A_STRIDE_VEC4, - gl_CooperativeMatrixLayoutRowMajor); - } - - coopmat matB; - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint col_b = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j) / FP16_PER_VEC4; - coopMatLoad( - matB, Bsh, - cur_base_B + k_start * B_STRIDE_VEC4 + col_b, - B_STRIDE_VEC4, - gl_CooperativeMatrixLayoutRowMajor); - - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - result[i][j] = coopMatMulAdd(matA[i], matB, result[i][j]); - } - } - } - - // --- store temp (chunk+1) -> nxt slice, dequantizing B --- - { - [[unroll]] for (uint p = 0; p < A_PASSES; ++p) { - Ash[nxt_base_A + (p * A_ROWS_PER_PASS + a_row_offset) * A_STRIDE_VEC4 + a_col] = - temp_A[p]; - } - [[unroll]] for (uint p = 0; p < B_PASSES; ++p) { -#ifdef WEIGHT_INT4 - Bsh[nxt_base_B + (p * B_ROWS_PER_PASS + b_row_offset) * B_STRIDE_VEC4 + b_col] = - dequant_block(temp_B[p], col_lo, col_hi, sc0, sc1); -#else - Bsh[nxt_base_B + (p * B_ROWS_PER_PASS + b_row_offset) * B_STRIDE_VEC4 + b_col] = - dequant_block(temp_Ba[p], temp_Bb[p], b_shift, sc0, sc1); -#endif - } - } - } - - // --- exit from MAIN LOOP: math on the last chunk --- - { - const uint cur_base_A = (chunk % 2u) * ASH_SLICE; - const uint cur_base_B = (chunk % 2u) * BSH_SLICE; - - barrier(); - - [[unroll]] for (uint k = 0; k < WG_TILE_K / MMA_K; ++k) { - const uint k_start = MMA_K * k; - - coopmat matA[MMAS_PER_SG_M]; - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - const uint row_a = MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); - coopMatLoad( - matA[i], Ash, - cur_base_A + row_a * A_STRIDE_VEC4 + k_start / FP16_PER_VEC4, - A_STRIDE_VEC4, - gl_CooperativeMatrixLayoutRowMajor); - } - - coopmat matB; - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint col_b = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j) / FP16_PER_VEC4; - coopMatLoad( - matB, Bsh, - cur_base_B + k_start * B_STRIDE_VEC4 + col_b, - B_STRIDE_VEC4, - gl_CooperativeMatrixLayoutRowMajor); - - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - result[i][j] = coopMatMulAdd(matA[i], matB, result[i][j]); - } - } - } - } - - // --- Bias staging (if any) --- -#ifdef HAS_BIAS - if (apply_bias > 0) { - for (uint t = gl_LocalInvocationID.x; t < WG_TILE_N; t += WG_SIZE) { - bias_sh[t] = float(t_bias[tile_n_start + t]); - } - memoryBarrierShared(); - barrier(); - } -#endif - - // --- Store result tile --- - // N for the store address math MUST come from the spec constant, not the - // sizes UBO (see out_N_arg above). - const uint N_out = uint(out_N_arg); - [[unroll]] for (uint i = 0; i < MMAS_PER_SG_M; ++i) { - [[unroll]] for (uint j = 0; j < MMAS_PER_SG_N; ++j) { - const uint gi = tile_m_start + MMA_M * (MMAS_PER_SG_M * warpInTile.y + i); - const uint gj = tile_n_start + MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - -#ifdef HAS_BIAS - if (apply_bias > 0) { - const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopmat bias_tile; - coopMatLoad( - bias_tile, bias_sh, - local_n, /*stride=*/0u, - gl_CooperativeMatrixLayoutRowMajor); - result[i][j] += bias_tile; - } -#endif - - coopmat out_tile = - coopmat(result[i][j]); - coopMatStore( - out_tile, t_output, - gi * N_out + gj, N_out, - gl_CooperativeMatrixLayoutRowMajor); - } - } -} diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_qw_coopmat.yaml b/backends/vulkan/runtime/graph/ops/glsl/linear_qw_coopmat.yaml deleted file mode 100644 index cd17e9edd6c..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_qw_coopmat.yaml +++ /dev/null @@ -1,47 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -# coopmat variant of the weight-only int4 quantized linear tiled shader (fp16 -# act x INT4 weight, dequantized to fp16 at the shared-memory store). -# WEIGHT_NBITS=4 -> linear_q4gsw_coopmat (INT4 group-symmetric) -# Forces buffer storage for activation/output (coopMatLoad/Store on buffers); -# INT weight storage can be texture2d or buffer (matches the tiled path). -# DTYPE = half only; fp32 activations are not supported. -# Geometry follows the double-buffered reference (coopmat_mm_ref): 128x128 -# tile, K-step 16, 4 subgroups x 32 threads (subgroup size 32 forced). The -# 128x128 / 2x2-subgroup-grid geometry is specs/036-portable-device-sweep's -# e2e-ranked winner on M51 (samsung xclipse 970), re-verified through this -# shipped shader path (not just the tsweep_ toggle) 2026-07-24 — confirmed -# +6.8%/+7.7%/+10.1% (1B/3B/8B prefill tok/s) over the prior 128x64/2x2 tile -# (kept as tsweep seed tsweep_t128x64k16g22s32 for future sweeps). NOTE: an -# EARLIER 128x128 shape (4x2 subgroup grid, not 2x2) was -25% vs 128x64/2x2 -# on the same device — don't conflate the two; the subgroup grid, not just -# the tile size, is what makes this one different. NOTE: the C++ dispatch in -# QuantizedLinear.cpp must keep kQ4gswCoopmatDims.n and .wg_size in sync with -# WG_TILE_N (128) and WG_SIZE (= SG_GRID_X*SG_GRID_Y*SUBGROUP = 128). - -linear_qw_coopmat: - parameter_names_with_default_values: - PRECISION: highp - HAS_BIAS: false - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - MMA_M: 16 - MMA_N: 16 - MMA_K: 16 - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 2 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - shader_variants: - - NAME: linear_q4gsw_coopmat_buffer_texture2d_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: texture2d - - NAME: linear_q4gsw_coopmat_buffer_buffer_half - WEIGHT_NBITS: 4 - WEIGHT_STORAGE: buffer From 6989c5d87c50fbcbc6851536f652772f62b2e723 Mon Sep 17 00:00:00 2001 From: Yanwen Xu Date: Tue, 15 Sep 2026 12:28:55 -0700 Subject: [PATCH 28/28] [ET-VK] Housekeeping pass 2 (cont'd): the two files a bad git-add dropped Companion to 1731d590ba, which should have included these two but a fatal git-add error on an already-staged path silently excluded them from that commit: - QuantizedLinear.cpp: the comment/constant fixes for the now-deleted linear_qw_coopmat.yaml / linear_dq8ca_qw_coopmat.yaml citations, per 1731d590ba's own message. - linear_q4gsw_coopmat_tsweep_dbuf4.yaml: the 97->3 shader_variants trim, same commit. No content difference from what was already described and rebuild-verified in 1731d590ba -- this just lands the two files that commit's message says are there but isn't. --- .../linear_q4gsw_coopmat_tsweep_dbuf4.yaml | 821 +----------------- .../graph/ops/impl/QuantizedLinear.cpp | 14 +- 2 files changed, 22 insertions(+), 813 deletions(-) diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_coopmat_tsweep_dbuf4.yaml b/backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_coopmat_tsweep_dbuf4.yaml index fbc01bc546f..2bd92ededc9 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_coopmat_tsweep_dbuf4.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_coopmat_tsweep_dbuf4.yaml @@ -4,17 +4,19 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -# specs/041-dbuf4-tile-sweep: TILE-SIZE SWEEP variants of the fp16 q4gsw -# coopmat kernel's dbuf4 ("store-first", prefetch-peeled) loop structure -- -# linear_q4gsw_coopmat_tsweep_dbuf4.glsl is a fork of -# linear_q4gsw_coopmat_tsweep.glsl (dbuf1) with only the loop structure -# swapped. Only the tile geometry (WG_TILE_*, SG_GRID_*, SUBGROUP_SIZE) -# varies per variant. Selected at dispatch via +# linear_q4gsw_coopmat_tsweep_dbuf4.glsl is the fp16 q4gsw coopmat kernel's +# dbuf4 ("store-first", prefetch-peeled) loop structure. Tile geometry +# (WG_TILE_*, SG_GRID_*, SUBGROUP_SIZE) is selectable at dispatch via # ET_VK_Q4GSW_COOPMAT_VARIANT=tsweep_dbuf4_txkgs -# (QuantizedLinear.cpp). Seed variant below matches the current production -# 4w tile (dbuf1, specs/036) as a legal, known-fast starting point for -# sweep.py's Optuna search; specs/041's sweep appends further candidates -# here as it runs. +# (QuantizedLinear.cpp), but t128x128k16g22s32 below is the only tile this +# branch ships or builds -- it is the current production 4w default. +# +# 2026-09-15 housekeeping: this file carried the full specs/041-dbuf4- +# tile-sweep search space (97 variants across ~48 tile geometries) even +# though only this one tile is ever dispatched. Trimmed to the 3 storage +# combos of the shipped tile. Full sweep history (every candidate tried, +# measured deltas) preserved on branch +# yanwen/release14-quant-shaders-archived-2026-09-15. linear_q4gsw_coopmat_tsweep_dbuf4: parameter_names_with_default_values: @@ -57,802 +59,3 @@ linear_q4gsw_coopmat_tsweep_dbuf4: SG_GRID_X: 2 SG_GRID_Y: 2 SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x128k16g41s32_buffer_texture2d_half - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 4 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x128k16g41s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 4 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t256x256k16g14s32_buffer_texture2d_half - WEIGHT_STORAGE: texture2d - WG_TILE_M: 256 - WG_TILE_N: 256 - WG_TILE_K: 16 - SG_GRID_X: 1 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t256x256k16g14s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_STORAGE: texture2d - WG_TILE_M: 256 - WG_TILE_N: 256 - WG_TILE_K: 16 - SG_GRID_X: 1 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x256k32g81s32_buffer_texture2d_half - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 256 - WG_TILE_K: 32 - SG_GRID_X: 8 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x256k32g81s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 256 - WG_TILE_K: 32 - SG_GRID_X: 8 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t32x64k64g41s64_buffer_texture2d_half - WEIGHT_STORAGE: texture2d - WG_TILE_M: 32 - WG_TILE_N: 64 - WG_TILE_K: 64 - SG_GRID_X: 4 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 64 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t32x64k64g41s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_STORAGE: texture2d - WG_TILE_M: 32 - WG_TILE_N: 64 - WG_TILE_K: 64 - SG_GRID_X: 4 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 64 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x128k64g14s32_buffer_texture2d_half - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 1 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x128k64g14s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 1 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t128x256k32g18s32_buffer_texture2d_half - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 256 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 8 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t128x256k32g18s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 256 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 8 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t16x64k64g11s32_buffer_texture2d_half - WEIGHT_STORAGE: texture2d - WG_TILE_M: 16 - WG_TILE_N: 64 - WG_TILE_K: 64 - SG_GRID_X: 1 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t16x64k64g11s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_STORAGE: texture2d - WG_TILE_M: 16 - WG_TILE_N: 64 - WG_TILE_K: 64 - SG_GRID_X: 1 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t256x128k32g44s32_buffer_texture2d_half - WEIGHT_STORAGE: texture2d - WG_TILE_M: 256 - WG_TILE_N: 128 - WG_TILE_K: 32 - SG_GRID_X: 4 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t256x128k32g44s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_STORAGE: texture2d - WG_TILE_M: 256 - WG_TILE_N: 128 - WG_TILE_K: 32 - SG_GRID_X: 4 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x128k16g22s32_buffer_texture2d_half - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 2 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x128k16g22s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 2 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t32x64k128g12s32_buffer_texture2d_half - WEIGHT_STORAGE: texture2d - WG_TILE_M: 32 - WG_TILE_N: 64 - WG_TILE_K: 128 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t32x64k128g12s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_STORAGE: texture2d - WG_TILE_M: 32 - WG_TILE_N: 64 - WG_TILE_K: 128 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x256k16g12s32_buffer_texture2d_half - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 256 - WG_TILE_K: 16 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x256k16g12s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 256 - WG_TILE_K: 16 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t16x64k64g21s64_buffer_texture2d_half - WEIGHT_STORAGE: texture2d - WG_TILE_M: 16 - WG_TILE_N: 64 - WG_TILE_K: 64 - SG_GRID_X: 2 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 64 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t16x64k64g21s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_STORAGE: texture2d - WG_TILE_M: 16 - WG_TILE_N: 64 - WG_TILE_K: 64 - SG_GRID_X: 2 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 64 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x32k64g22s64_buffer_texture2d_half - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 32 - WG_TILE_K: 64 - SG_GRID_X: 2 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x32k64g22s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 32 - WG_TILE_K: 64 - SG_GRID_X: 2 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t128x256k16g81s32_buffer_texture2d_half - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 256 - WG_TILE_K: 16 - SG_GRID_X: 8 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t128x256k16g81s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 256 - WG_TILE_K: 16 - SG_GRID_X: 8 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t16x32k128g21s64_buffer_texture2d_half - WEIGHT_STORAGE: texture2d - WG_TILE_M: 16 - WG_TILE_N: 32 - WG_TILE_K: 128 - SG_GRID_X: 2 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 64 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t16x32k128g21s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_STORAGE: texture2d - WG_TILE_M: 16 - WG_TILE_N: 32 - WG_TILE_K: 128 - SG_GRID_X: 2 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 64 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x128k64g22s64_buffer_texture2d_half - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 2 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x128k64g22s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 2 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x128k64g81s32_buffer_texture2d_half - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 8 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x128k64g81s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 8 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t32x128k64g41s32_buffer_texture2d_half - WEIGHT_STORAGE: texture2d - WG_TILE_M: 32 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 4 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t32x128k64g41s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_STORAGE: texture2d - WG_TILE_M: 32 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 4 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t16x64k32g11s64_buffer_texture2d_half - WEIGHT_STORAGE: texture2d - WG_TILE_M: 16 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 64 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t16x64k32g11s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_STORAGE: texture2d - WG_TILE_M: 16 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 64 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t16x16k64g11s64_buffer_texture2d_half - WEIGHT_STORAGE: texture2d - WG_TILE_M: 16 - WG_TILE_N: 16 - WG_TILE_K: 64 - SG_GRID_X: 1 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 64 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t16x16k64g11s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_STORAGE: texture2d - WG_TILE_M: 16 - WG_TILE_N: 16 - WG_TILE_K: 64 - SG_GRID_X: 1 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 64 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t128x64k16g22s32_buffer_texture2d_half - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 16 - SG_GRID_X: 2 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t128x64k16g22s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 16 - SG_GRID_X: 2 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t32x128k32g22s32_buffer_texture2d_half - WEIGHT_STORAGE: texture2d - WG_TILE_M: 32 - WG_TILE_N: 128 - WG_TILE_K: 32 - SG_GRID_X: 2 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t32x128k32g22s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_STORAGE: texture2d - WG_TILE_M: 32 - WG_TILE_N: 128 - WG_TILE_K: 32 - SG_GRID_X: 2 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x128k32g11s32_buffer_texture2d_half - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x128k32g11s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 32 - SG_GRID_X: 1 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t128x64k64g12s32_buffer_texture2d_half - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 64 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t128x64k64g12s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 64 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x128k16g12s64_buffer_texture2d_half - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x128k16g12s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t128x128k16g21s32_buffer_texture2d_half - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 2 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t128x128k16g21s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 2 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t16x64k32g21s32_buffer_texture2d_half - WEIGHT_STORAGE: texture2d - WG_TILE_M: 16 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 2 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t16x64k32g21s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_STORAGE: texture2d - WG_TILE_M: 16 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 2 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x128k16g14s32_buffer_texture2d_half - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 1 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x128k16g14s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 1 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t16x32k64g11s32_buffer_texture2d_half - WEIGHT_STORAGE: texture2d - WG_TILE_M: 16 - WG_TILE_N: 32 - WG_TILE_K: 64 - SG_GRID_X: 1 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t16x32k64g11s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_STORAGE: texture2d - WG_TILE_M: 16 - WG_TILE_N: 32 - WG_TILE_K: 64 - SG_GRID_X: 1 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x32k64g21s64_buffer_texture2d_half - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 32 - WG_TILE_K: 64 - SG_GRID_X: 2 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 64 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x32k64g21s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 32 - WG_TILE_K: 64 - SG_GRID_X: 2 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 64 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t32x64k64g11s32_buffer_texture2d_half - WEIGHT_STORAGE: texture2d - WG_TILE_M: 32 - WG_TILE_N: 64 - WG_TILE_K: 64 - SG_GRID_X: 1 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t32x64k64g11s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_STORAGE: texture2d - WG_TILE_M: 32 - WG_TILE_N: 64 - WG_TILE_K: 64 - SG_GRID_X: 1 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t128x64k16g21s32_buffer_texture2d_half - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 16 - SG_GRID_X: 2 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t128x64k16g21s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 16 - SG_GRID_X: 2 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x128k64g21s64_buffer_texture2d_half - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 2 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 64 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x128k64g21s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 2 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 64 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x128k32g22s64_buffer_texture2d_half - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 32 - SG_GRID_X: 2 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x128k32g22s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 32 - SG_GRID_X: 2 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t32x128k64g22s64_buffer_texture2d_half - WEIGHT_STORAGE: texture2d - WG_TILE_M: 32 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 2 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t32x128k64g22s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_STORAGE: texture2d - WG_TILE_M: 32 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 2 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 64 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t128x128k16g14s32_buffer_texture2d_half - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 1 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t128x128k16g14s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 1 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t128x64k16g12s32_buffer_texture2d_half - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 16 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t128x64k16g12s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 16 - SG_GRID_X: 1 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x128k16g21s32_buffer_texture2d_half - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 2 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x128k16g21s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 16 - SG_GRID_X: 2 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t128x64k32g24s32_buffer_texture2d_half - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t128x64k32g24s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 2 - SG_GRID_Y: 4 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t32x64k32g22s32_buffer_texture2d_half - WEIGHT_STORAGE: texture2d - WG_TILE_M: 32 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 2 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t32x64k32g22s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_STORAGE: texture2d - WG_TILE_M: 32 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 2 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t32x256k32g21s32_buffer_texture2d_half - WEIGHT_STORAGE: texture2d - WG_TILE_M: 32 - WG_TILE_N: 256 - WG_TILE_K: 32 - SG_GRID_X: 2 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t32x256k32g21s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_STORAGE: texture2d - WG_TILE_M: 32 - WG_TILE_N: 256 - WG_TILE_K: 32 - SG_GRID_X: 2 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x128k64g41s32_buffer_texture2d_half - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 4 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x128k64g41s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 128 - WG_TILE_K: 64 - SG_GRID_X: 4 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t128x32k64g22s32_buffer_texture2d_half - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 32 - WG_TILE_K: 64 - SG_GRID_X: 2 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t128x32k64g22s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 32 - WG_TILE_K: 64 - SG_GRID_X: 2 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t16x64k128g21s32_buffer_texture2d_half - WEIGHT_STORAGE: texture2d - WG_TILE_M: 16 - WG_TILE_N: 64 - WG_TILE_K: 128 - SG_GRID_X: 2 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t16x64k128g21s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_STORAGE: texture2d - WG_TILE_M: 16 - WG_TILE_N: 64 - WG_TILE_K: 128 - SG_GRID_X: 2 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t128x64k64g42s32_buffer_texture2d_half - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 64 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t128x64k64g42s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 64 - WG_TILE_K: 64 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t128x128k32g42s32_buffer_texture2d_half - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 32 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t128x128k32g42s32_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_STORAGE: texture2d - WG_TILE_M: 128 - WG_TILE_N: 128 - WG_TILE_K: 32 - SG_GRID_X: 4 - SG_GRID_Y: 2 - SUBGROUP_SIZE: 32 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x64k64g11s64_buffer_texture2d_half - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 64 - WG_TILE_K: 64 - SG_GRID_X: 1 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 64 - - NAME: linear_q4gsw_coopmat_tsweep_dbuf4_t64x64k64g11s64_texture3d_texture2d_half - IO_STORAGE: texture3d - WEIGHT_STORAGE: texture2d - WG_TILE_M: 64 - WG_TILE_N: 64 - WG_TILE_K: 64 - SG_GRID_X: 1 - SG_GRID_Y: 1 - SUBGROUP_SIZE: 64 diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp index 001a94b236e..a3facad210a 100644 --- a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp @@ -76,11 +76,17 @@ struct CoopmatTileDims { // never reads it. 0 = "unknown / shipped default". uint32_t sg_grid_y; }; -// linear_qw_coopmat.yaml: 128x128, 2x2 subgroup grid, sg32 -> WG_SIZE 128. +// Defensive-only fallback for parse_tsweep_tile(): both +// {q4gsw,dq8ca}_coopmat_variant() are guaranteed by construction to always +// return a prefix-whitelisted string (the hardcoded default, or an +// env-var value already validated by is_*_shippable_token), so +// parse_tsweep_tile always parses the real dims out of that string and +// this fallback is never actually returned in current code. Kept at the +// current shipped tile's own dims (not some other tile) purely so that if +// the invariant above is ever violated, the failure mode is "silently use +// today's real default" rather than a stale, unrelated geometry. constexpr CoopmatTileDims kQ4gswCoopmatDims = {128, 128, 16, 128, 2}; -// linear_dq8ca_qw_coopmat.yaml: 64x32, 1x2 grid, sg64 -> WG_SIZE 128 -// (specs/027-e2e-tile-sweep winner, was 128x64x32/256). -constexpr CoopmatTileDims kDq8caQ4gswCoopmatDims = {64, 32, 32, 128, 2}; +constexpr CoopmatTileDims kDq8caQ4gswCoopmatDims = {128, 64, 32, 256, 2}; // specs/028-4w-e2e-tile-sweep / specs/041-dbuf4-tile-sweep: // ET_VK_Q4GSW_COOPMAT_VARIANT / ET_VK_DQ8CA_COOPMAT_VARIANT can swap the