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_ = 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..97eed7ddf67 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_coopmat_tsweep_dbuf4zpgtr.glsl @@ -0,0 +1,766 @@ +/* + * 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; + + // 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) { + 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. + // 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_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/glsl/linear_dq8ca_qw_coopmat.glsl b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_qw_coopmat.glsl deleted file mode 100644 index 755261452f4..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_qw_coopmat.glsl +++ /dev/null @@ -1,588 +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 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 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. - * - * 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 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). - * - * 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. -${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); - } - - // Prefetch chunk 0 into temp registers, then store to slice 0 (no barrier; - // the first loop iteration's barrier publishes it). - 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 — 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. - // 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). - 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 959cb51966d..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_qw_coopmat.yaml +++ /dev/null @@ -1,40 +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. -# 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). - -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: 128 - WG_TILE_N: 64 - WG_TILE_K: 32 - SG_GRID_X: 2 - 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_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/linear_qw_coopmat.glsl b/backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_coopmat_tsweep_dbuf4.glsl similarity index 55% rename from backends/vulkan/runtime/graph/ops/glsl/linear_qw_coopmat.glsl rename to backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_coopmat_tsweep_dbuf4.glsl index 1f9707c3b79..08447b06059 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_qw_coopmat.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_coopmat_tsweep_dbuf4.glsl @@ -7,57 +7,36 @@ */ /* - * 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) + * 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). * - * 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. + * 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 - * INT4: group_size % WG_TILE_K == 0 (each group = whole number of chunks) - * Misaligned shapes silently miscompute / overrun — gate at dispatch time. + * 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. */ -// REQUIRED_SUBGROUP_SIZE = 32 - #version 450 core #extension GL_KHR_cooperative_matrix : require @@ -69,23 +48,21 @@ #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 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, "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)} @@ -96,20 +73,11 @@ ${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) --- +// --- 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}; @@ -142,7 +110,17 @@ const uint BSH_SLICE = WG_TILE_K * B_STRIDE_VEC4; shared uvec4 Ash[2 * ASH_SLICE]; shared uvec4 Bsh[2 * BSH_SLICE]; #ifdef HAS_BIAS -shared float bias_sh[WG_TILE_N]; +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. @@ -153,12 +131,10 @@ 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 +// FP16 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( @@ -167,48 +143,39 @@ uvec4 dequant_block( 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; + 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)); } -#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; +// 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 } -#endif // WEIGHT_INT4 - void main() { const uvec2 tileID = uvec2(gl_WorkGroupID.xy); const uvec2 warpInTile = uvec2( @@ -219,20 +186,15 @@ void main() { 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); + result[i][j] = coopmat(0.0); } } @@ -241,103 +203,56 @@ void main() { 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 + ivec4 temp_B[B_PASSES]; // ========================================================= - // PROLOGUE: prefetch chunk 0 into temp registers, then store to slice 0. + // 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; - 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 + 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]; -#else [[unroll]] for (uint p = 0; p < B_PASSES; ++p) { - const uint k4 = (p * B_ROWS_PER_PASS + b_row_offset) >> 2u; + const uint k_row = p * B_ROWS_PER_PASS + b_row_offset; + ivec4 wblock; #ifdef WEIGHT_BUFFER - temp_Ba[p] = t_packed_weight[k4 * N4 + n4_a]; - temp_Bb[p] = t_packed_weight[k4 * N4 + n4_a + 1u]; + wblock = t_packed_weight[n8_blk * K4 + (k_row >> 2u)]; #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); + wblock = texelFetch(t_packed_weight, ivec2(k_row >> 2u, n8_blk), 0); #endif + temp_B[p] = wblock; } -#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: + // 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) @@ -350,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 --- @@ -359,13 +281,8 @@ void main() { [[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)); + temp_A[p] = load_a_vec4(row, k_hv4, K4); } -#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 @@ -380,18 +297,6 @@ void main() { 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 --- @@ -430,22 +335,24 @@ void main() { 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 --- + // --- 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; + // 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) { @@ -481,7 +388,7 @@ void main() { #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]); + bias_sh[t] = float16_t(t_bias[tile_n_start + t]); } memoryBarrierShared(); barrier(); @@ -489,8 +396,74 @@ void main() { #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 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. + // 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, /*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) { @@ -500,7 +473,7 @@ void main() { #ifdef HAS_BIAS if (apply_bias > 0) { const uint local_n = MMA_N * (MMAS_PER_SG_N * warpInTile.x + j); - coopmat bias_tile; + coopmat bias_tile; coopMatLoad( bias_tile, bias_sh, local_n, /*stride=*/0u, @@ -517,4 +490,5 @@ void main() { 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..2bd92ededc9 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_coopmat_tsweep_dbuf4.yaml @@ -0,0 +1,61 @@ +# 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. + +# 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), 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: + 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 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 173caf8dc29..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_qw_coopmat.yaml +++ /dev/null @@ -1,40 +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): 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 -# 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). - -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: 64 - 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 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/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/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_attn_weights_coopmat.glsl b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_coopmat.glsl new file mode 100644 index 00000000000..8b4e95815ad --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_coopmat.glsl @@ -0,0 +1,383 @@ +/* + * 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). + * + * 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. + */ + +#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")} +// 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}; +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); + + // 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) { + 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; + } + + // 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) --- + [[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]); + } + } + } + + // 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(); + } + + 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) { + [[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); + } + } + // 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 --- + 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..c40737e0587 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_coopmat.yaml @@ -0,0 +1,49 @@ +# 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 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: + 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_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_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/glsl/sdpa_compute_out_coopmat.glsl b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coopmat.glsl new file mode 100644 index 00000000000..d142608f7bd --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coopmat.glsl @@ -0,0 +1,283 @@ +/* + * 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); + + // 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, + // 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)); + } + + // 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) --- + [[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]); + } + } + } + + // 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(); + } + + // --- 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..b7859e358cd --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coopmat.yaml @@ -0,0 +1,46 @@ +# 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 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: + 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_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/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/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 b9a03d85161..a3facad210a 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,298 @@ 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; +}; +// 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}; +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 +// 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'). +// 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. +// +// 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_dbuf4_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[] = { + // 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", }; -// 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}; + +// (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 std::strlen(prefix); + } + } + 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() { + // 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("tsweep_dbuf4_t128x128k16g22s32"); + } + const std::string v(env); + 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_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 + // 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 + // (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 + // 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"); + if (!env) { + return std::string("tsweep_dbuf4zpgtr_t128x64k32g42s32"); + } + const std::string v(env); + if (is_dq8ca_shippable_token(v)) { + return v; + } + // 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; +} + +// Parses "tsweep_txkgs" or +// "tsweep_dbuf_txkgs" -> {M, N, K, SGX*SGY*sub, +// SGY}. Returns fallback unchanged if the token matches none of +// any family's prefix list. +static CoopmatTileDims parse_tsweep_tile( + const std::string& variant, + const CoopmatTileDims& fallback) { + 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; + } + 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 +437,36 @@ utils::uvec3 quantized_linear_local_wg_size( } } +// 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 = [] { + 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; +} + // 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 +479,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 +508,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); @@ -218,6 +582,68 @@ 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() { + 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_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 +// (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, @@ -236,18 +662,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)); @@ -294,26 +731,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)); - if (can_use_q4gsw_coopmat( - graph, - out, - fp_input, - group_size, - resize_args.at(2), - kDq8caQ4gswCoopmatDims.m, - kDq8caQ4gswCoopmatDims.n, - kDq8caQ4gswCoopmatDims.k)) { - std::string kernel_name = "linear_dq8ca_q4gsw_coopmat"; - 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"; @@ -781,13 +1217,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) { @@ -833,22 +1293,42 @@ 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); - 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, @@ -932,6 +1412,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) { @@ -970,6 +1488,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/graph/ops/impl/SDPA.cpp b/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp index 3efb834725d..54a4b5f67aa 100644 --- a/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp @@ -21,7 +21,10 @@ #include #include +#include #include +#include +#include namespace vkcompute { @@ -178,16 +181,193 @@ 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). +// +// 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; +} + +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, + 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) { + 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); + 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)); + 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 +394,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 +401,23 @@ 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 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), dims.n); + const uint32_t num_tiles_m = + utils::div_up(static_cast(d.S), dims.m); + return { + num_tiles_n * sdpa_wg_size(dims), + 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 +432,13 @@ 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 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 {sdpa_wg_size(sdpa_attn_tile_dims()), 1, 1}; + } const bool use_coop_algorithm = shader.kernel_name.find("_coop") != std::string::npos; if (use_coop_algorithm) { @@ -293,10 +496,34 @@ 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); + 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)); + 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 +546,27 @@ 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 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), dims.n); + const uint32_t num_tiles_m = + utils::div_up(static_cast(d.S), dims.m); + return { + num_tiles_n * sdpa_wg_size(dims), + 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 +580,13 @@ 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 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 {sdpa_wg_size(sdpa_out_tile_dims()), 1, 1}; + } const bool use_coop_algorithm = shader.kernel_name.find("_coop") != std::string::npos; if (use_coop_algorithm) { @@ -419,6 +667,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, @@ -430,8 +690,28 @@ 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), + // 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 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), + aw_row_width_at_construction}, // Resize Args: [q, k, input_pos_symint_or_dummy, mode] {q, k, input_pos_symint, mode_ref}, // Resizing Logic @@ -511,6 +791,22 @@ 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) 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 + 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( graph, pick_sdpa_av_shader, @@ -522,8 +818,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 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( 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/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..0c64bf669d2 --- /dev/null +++ b/backends/vulkan/test/custom_ops/test_llama_microbench.cpp @@ -0,0 +1,2582 @@ +// 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 +// --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 +// 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 +// +// 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 +#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; + // 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; +}; + +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 << "," + << 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; +} +// 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; + 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 (!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}; + // 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; +} + +// 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) { + 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]); + 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"; + } + } + 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; +} + +// --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; +}; +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]); + 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) + : -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 + 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 +}; + +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 softmax_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; + 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; + // 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) { + 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); + softmax_timings_us.push_back(softmax_time_us); + total_timings_us.push_back(qk_time_us + av_time_us + softmax_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.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; +} + +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}, + {"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}, + }; + 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; +} + +// ===================== 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 + // "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, "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, "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 +// (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. +// --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) { + if (g_sdpa_force_fallback) { + setenv("ET_VK_DISABLE_COOPMAT", "1", 1); + } else { + 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; + } + } + } + + // 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 + << " 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"; + unsetenv("ET_VK_DISABLE_COOPMAT"); // restore the tree's default-on state + 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(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; +} + +// ============================== 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. +// --- 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"; + 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" + " --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"; +} + +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, 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 + // 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]; + 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 == "--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") { + 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) { + 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") + << ",max_shared_mem_bytes=" + << adapter->max_compute_shared_memory_size() << "\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. + 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 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(sdpa_tier.c_str())); + } + 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; + } + } + + 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); + + 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 + // 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; +} 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.