mirror of
https://github.com/leejet/stable-diffusion.cpp.git
synced 2026-11-11 03:24:30 +00:00
perf(vulkan): accelerate packed INT4 with cooperative matrices (#2280)
This commit is contained in:
leejet authored and GitHub committed
1 parent
f5df5f8b84
commit
5a1ba5d895
6 files changed
+420
-6
No files matched your search
+20
-3
@@ -133,9 +133,26 @@ Other CUDA shapes and builds retain the packed byte-dot path. Set
|
||||
use packed byte-dot instructions, with portable byte arithmetic for other targets.
|
||||
Vulkan requires integer dot product support and at least 13312 bytes of shared
|
||||
memory for W4A4 or 16452 bytes for W4A8. GPU W4 kernels require the logical input
|
||||
width to be divisible by 64; H64 and H256 are supported. Other shapes use CPU
|
||||
fallback. HIP W4A8 selects wave64 WMMA by the active compilation target's
|
||||
intrinsic support, preferring the GFX12 fragment layout over the GFX11 layout.
|
||||
width to be divisible by 64; H64 and H256 are supported. When supported by the
|
||||
compiler and device, Vulkan uses NV cooperative matrix integer GEMM. W4A4
|
||||
decodes packed weights directly into cooperative matrices for aligned input
|
||||
and output widths. Complete activation tiles and the final partial tile use
|
||||
separate shader specializations. Other W4A4 shapes use double-buffered weight
|
||||
and activation tiles (20480 bytes of shared memory, or 30720 for larger tiles).
|
||||
W4A8 stages both weights and activations in shared memory in K=64 tiles,
|
||||
reusing them for each 16-value group's integer
|
||||
MMA and FP32 scale restoration. FP8 scales use a shared lookup table, and
|
||||
expanding projections visit activation tiles in groups of weight panels for
|
||||
cache reuse. Adjacent groups share a K=32 activation fragment, with
|
||||
complementary zero padding in the weights to preserve each group's scale.
|
||||
This path uses 17476 bytes of shared memory, or 22596 for larger tiles.
|
||||
Both paths use 256-thread workgroups and keep the model's
|
||||
weights packed; neither expands a full layer in device memory. Other Vulkan
|
||||
devices retain the packed dot-product path; set `GGML_VK_W4_FORCE_DP4A=1` to compare
|
||||
against it. FP32 accumulation order can differ between these paths.
|
||||
Other shapes use CPU fallback. HIP W4A8 selects wave64 WMMA by the active
|
||||
compilation target's intrinsic support, preferring the GFX12 fragment layout
|
||||
over the GFX11 layout.
|
||||
There is no GPU architecture allowlist: newer targets supporting either
|
||||
intrinsic automatically use its matching layout. Selection happens per device
|
||||
image, including multi-architecture builds; targets supporting neither intrinsic
|
||||
|
||||
@@ -56,6 +56,9 @@ sd_vulkan_shader(quantize_i8_convrot quantize_i8_convrot ${SD_VK_DOT} -O)
|
||||
sd_vulkan_shader(quantize_i8_convrot_h64 quantize_i8_convrot_h64 ${SD_VK_DOT} -O -DGROUP_SIZE=64)
|
||||
sd_vulkan_shader(mul_mat_w4a4_convrot mul_mat_w4a4_convrot ${SD_VK_DOT} -O)
|
||||
sd_vulkan_shader(mul_mat_w4a8_convrot mul_mat_w4a8_convrot ${SD_VK_DOT} -O)
|
||||
sd_vulkan_shader(mul_mat_w4a4_convrot_cm2 mul_mat_w4a4_convrot_cm2 ${SD_VK_CM2} -O)
|
||||
sd_vulkan_shader(mul_mat_w4a4_convrot_decode_cm2 mul_mat_w4a4_convrot_decode_cm2 ${SD_VK_CM2} -O)
|
||||
sd_vulkan_shader(mul_mat_w4a8_convrot_cm2 mul_mat_w4a8_convrot_cm2 ${SD_VK_CM2} -O)
|
||||
sd_vulkan_shader(quantize_i8_convrot_subgroup quantize_i8_convrot_subgroup ${SD_VK_DOT} -O)
|
||||
sd_vulkan_shader(quantize_i8_convrot_register quantize_i8_convrot_subgroup ${SD_VK_DOT} -O -DSINGLE_PASS=1)
|
||||
sd_vulkan_shader(swiglu_quantize_i8_convrot quantize_i8_convrot_subgroup ${SD_VK_DOT} -O -DSWIGLU=1)
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
#include <algorithm>
|
||||
#include <array>
|
||||
#include <climits>
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
#include <vector>
|
||||
#include "core/ggml_extensions.h"
|
||||
@@ -12,7 +13,10 @@
|
||||
#include "mul_mat_i8_tensorwise_cm1.h"
|
||||
#include "mul_mat_i8_tensorwise_cm2.h"
|
||||
#include "mul_mat_w4a4_convrot.h"
|
||||
#include "mul_mat_w4a4_convrot_cm2.h"
|
||||
#include "mul_mat_w4a4_convrot_decode_cm2.h"
|
||||
#include "mul_mat_w4a8_convrot.h"
|
||||
#include "mul_mat_w4a8_convrot_cm2.h"
|
||||
#include "quantize_i8_convrot.h"
|
||||
#include "quantize_i8_convrot_h64.h"
|
||||
#include "quantize_i8_convrot_register.h"
|
||||
@@ -23,11 +27,17 @@
|
||||
#include "swiglu_quantize_i8_convrot_16384.h"
|
||||
|
||||
namespace sd_vulkan_i8 {
|
||||
// Keep these budgets in sync with the cooperative matrix shader tile layouts.
|
||||
static constexpr uint32_t w4a4_shared = 2 * (64 + 64) * (64 / 4 + 4) * 4;
|
||||
static constexpr uint32_t w4a4_large_shared = 2 * (64 + 128) * (64 / 4 + 4) * 4;
|
||||
static constexpr uint32_t w4a8_shared = (64 * (64 / 2 + 4) + 64 * (64 / 16) + 64 * (64 / 4 + 4) + 256 + 256 + 16 + 1) * 4;
|
||||
static constexpr uint32_t w4a8_large_shared = (64 * (64 / 2 + 4) + 64 * (64 / 16) + 128 * (64 / 4 + 4) + 256 + 256 + 16 + 1) * 4;
|
||||
struct Device {
|
||||
const ggml_ext_vulkan_api* api;
|
||||
ggml_ext_vulkan_device_info info;
|
||||
ggml_ext_vulkan_compute_info compute;
|
||||
bool subgroup, cm2, cm2_small;
|
||||
bool w4a4_cm2, w4a8_cm2;
|
||||
uint32_t cm_m, cm_n, cm_k;
|
||||
};
|
||||
static bool device_info(ggml_backend_dev_t dev, Device& d) {
|
||||
@@ -55,6 +65,10 @@ namespace sd_vulkan_i8 {
|
||||
d.cm2 = true;
|
||||
d.cm2_small |= 64 % prop.m == 0 && 32 % prop.n == 0;
|
||||
}
|
||||
if (prop.flexible && prop.workgroup_scope && prop.invocations == 256 && d.info.max_group_size[0] >= 256 && 64 % prop.m == 0 && 64 % prop.n == 0) {
|
||||
d.w4a4_cm2 |= d.compute.shared_memory >= w4a4_shared && 64 % prop.k == 0 && sizeof(sd_vulkan_shaders::mul_mat_w4a4_convrot_cm2) > 4;
|
||||
d.w4a8_cm2 |= d.compute.shared_memory >= w4a8_shared && 32 % prop.k == 0 && sizeof(sd_vulkan_shaders::mul_mat_w4a8_convrot_cm2) > 4;
|
||||
}
|
||||
}
|
||||
return d.info.integer_dot_product && sizeof(sd_vulkan_shaders::mul_mat_i8_tensorwise) > 4;
|
||||
}
|
||||
@@ -698,9 +712,45 @@ namespace sd_vulkan_i8 {
|
||||
struct Parameters {
|
||||
uint32_t n, rows, k, bias;
|
||||
};
|
||||
const Parameters pc = {n, rows, k, op.bias ? 1u : 0u};
|
||||
auto pipeline = op.kind == SDW4Kind::W4A4 ? SD_VK_PIPELINE(mul_mat_w4a4_convrot, 6, Parameters) : SD_VK_PIPELINE(mul_mat_w4a8_convrot, 8, Parameters);
|
||||
const uint32_t groups[] = {(n + 63) / 64, (rows + 31) / 32, 1};
|
||||
const Parameters pc = {n, rows, k, op.bias ? 1u : 0u};
|
||||
auto pipeline = op.kind == SDW4Kind::W4A4 ? SD_VK_PIPELINE(mul_mat_w4a4_convrot, 6, Parameters) : SD_VK_PIPELINE(mul_mat_w4a8_convrot, 8, Parameters);
|
||||
uint32_t tile_cols = 64, tile_rows = 32;
|
||||
const bool large_w4a4 = op.kind == SDW4Kind::W4A4 && rows >= 512 && d.compute.shared_memory >= w4a4_large_shared;
|
||||
const bool large_w4a8 = op.kind == SDW4Kind::W4A8 && rows >= 512 && d.compute.shared_memory >= w4a8_large_shared;
|
||||
const bool decode_w4a4 = op.kind == SDW4Kind::W4A4 && d.w4a4_cm2 && k % 256 == 0 && n % 64 == 0 && rows >= 128 &&
|
||||
aligned(op.weight, 16) && sizeof(sd_vulkan_shaders::mul_mat_w4a4_convrot_decode_cm2) > 4;
|
||||
const uint32_t constants[] = {64, large_w4a4 || large_w4a8 ? 128u : 64u, 64};
|
||||
const char* force_dp4a = std::getenv("GGML_VK_W4_FORCE_DP4A");
|
||||
if (!(force_dp4a && std::atoi(force_dp4a)) && aligned(op.input, 16) &&
|
||||
(op.kind == SDW4Kind::W4A4 ? d.w4a4_cm2 : d.w4a8_cm2 && aligned(op.weight, 8))) {
|
||||
pipeline = op.kind == SDW4Kind::W4A4 ? SD_VK_PIPELINE(mul_mat_w4a4_convrot_cm2, 6, Parameters) : SD_VK_PIPELINE(mul_mat_w4a8_convrot_cm2, 8, Parameters);
|
||||
if (decode_w4a4) {
|
||||
struct DecodeParameters {
|
||||
uint32_t n, rows, k, bias, row_offset;
|
||||
};
|
||||
DecodeParameters decoded = {n, rows, k, pc.bias, 0};
|
||||
uint32_t decode_constants[] = {64, 128, 128, 1};
|
||||
pipeline = SD_VK_PIPELINE(mul_mat_w4a4_convrot_decode_cm2, 6, DecodeParameters);
|
||||
pipeline.constants = decode_constants;
|
||||
pipeline.constant_count = 4;
|
||||
uint32_t groups[] = {n / 64, rows / 128, 1};
|
||||
auto status = d.api->dispatch(e->backend, &pipeline, bindings.data(), &decoded, groups);
|
||||
if (status == GGML_STATUS_SUCCESS && rows % 128) {
|
||||
// Specialize the tail separately so bounds checks do not inhibit full-tile loads.
|
||||
decoded.row_offset = rows / 128 * 128;
|
||||
decode_constants[1] = rows % 128 <= 64 ? 64 : 128;
|
||||
decode_constants[3] = 0;
|
||||
groups[1] = 1;
|
||||
status = d.api->dispatch(e->backend, &pipeline, bindings.data(), &decoded, groups);
|
||||
}
|
||||
return finish(e, d, output, dst, alias, status);
|
||||
}
|
||||
pipeline.constants = constants;
|
||||
pipeline.constant_count = 3;
|
||||
tile_cols = constants[0];
|
||||
tile_rows = constants[1];
|
||||
}
|
||||
const uint32_t groups[] = {(n + tile_cols - 1) / tile_cols, (rows + tile_rows - 1) / tile_rows, 1};
|
||||
auto status = d.api->dispatch(e->backend, &pipeline, bindings.data(), &pc, groups);
|
||||
return finish(e, d, output, dst, alias, status);
|
||||
}
|
||||
|
||||
@@ -0,0 +1,88 @@
|
||||
#version 460
|
||||
#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require
|
||||
#extension GL_EXT_shader_explicit_arithmetic_types_int32 : require
|
||||
#extension GL_KHR_cooperative_matrix : require
|
||||
#extension GL_NV_cooperative_matrix2 : require
|
||||
#extension GL_KHR_memory_scope_semantics : require
|
||||
layout(local_size_x = 256) in;
|
||||
layout(constant_id = 0) const uint BM = 64;
|
||||
layout(constant_id = 1) const uint BN = 64;
|
||||
layout(constant_id = 2) const uint BK = 64;
|
||||
const uint STRIDE = BK / 4 + 4;
|
||||
layout(push_constant) uniform parameter {
|
||||
uint n;
|
||||
uint rows;
|
||||
uint k;
|
||||
uint has_bias;
|
||||
}
|
||||
p;
|
||||
layout(binding = 0, std430) readonly buffer W {
|
||||
uint data_w[];
|
||||
};
|
||||
layout(binding = 1, std430) readonly buffer A {
|
||||
uvec4 data_a[];
|
||||
};
|
||||
layout(binding = 2, std430) readonly buffer WS {
|
||||
float data_ws[];
|
||||
};
|
||||
layout(binding = 3, std430) readonly buffer AS {
|
||||
float data_as[];
|
||||
};
|
||||
layout(binding = 4, std430) readonly buffer B {
|
||||
float data_b[];
|
||||
};
|
||||
layout(binding = 5, std430) writeonly buffer D {
|
||||
int32_t data_d[];
|
||||
};
|
||||
shared uint weights[2 * BM * STRIDE];
|
||||
shared uint activations[2 * BN * STRIDE];
|
||||
void load_tile(uint base, uint buf, uint row_base, uint col_base) {
|
||||
const uint tid = gl_LocalInvocationIndex;
|
||||
for (uint i = tid; i < BN * BK / 16; i += 256) {
|
||||
const uint r = i / (BK / 16), t = i % (BK / 16);
|
||||
const uvec4 value = row_base + r < p.rows ? data_a[(row_base + r) * (p.k / 16) + base / 16 + t] : uvec4(0);
|
||||
const uint offset = (buf * BN + r) * STRIDE + t * 4;
|
||||
activations[offset] = value.x;
|
||||
activations[offset + 1] = value.y;
|
||||
activations[offset + 2] = value.z;
|
||||
activations[offset + 3] = value.w;
|
||||
}
|
||||
for (uint i = tid; i < BM * BK / 8; i += 256) {
|
||||
const uint c = i / (BK / 8), t = i % (BK / 8);
|
||||
const uint v = col_base + c < p.n ? data_w[(col_base + c) * (p.k / 8) + base / 8 + t] : 0u;
|
||||
const uint lo = (v & 15u) | ((v << 4) & 0xf00u) | ((v << 8) & 0xf0000u) | ((v << 12) & 0xf000000u);
|
||||
const uint hi = ((v >> 16) & 15u) | ((v >> 12) & 0xf00u) | ((v >> 8) & 0xf0000u) | ((v >> 4) & 0xf000000u);
|
||||
weights[(buf * BM + c) * STRIDE + 2 * t] = (lo & 0x07070707u) | ((lo & 0x08080808u) * 31u);
|
||||
weights[(buf * BM + c) * STRIDE + 2 * t + 1] = (hi & 0x07070707u) | ((hi & 0x08080808u) * 31u);
|
||||
}
|
||||
}
|
||||
int32_t dequantize(const in uint r, const in uint c, const in int32_t v, const in uint row, const in uint col) {
|
||||
if (row + c >= p.rows || col + r >= p.n)
|
||||
return 0;
|
||||
precise float value = float(v) * data_as[row + c] * data_ws[col + r];
|
||||
if (p.has_bias != 0)
|
||||
value += data_b[col + r];
|
||||
return floatBitsToInt(value);
|
||||
}
|
||||
void main() {
|
||||
const uint row = gl_WorkGroupID.y * BN, col = gl_WorkGroupID.x * BM;
|
||||
coopmat<int32_t, gl_ScopeWorkgroup, BM, BN, gl_MatrixUseAccumulator> sum = coopmat<int32_t, gl_ScopeWorkgroup, BM, BN, gl_MatrixUseAccumulator>(0);
|
||||
load_tile(0, 0, row, col);
|
||||
barrier();
|
||||
for (uint base = 0; base < p.k; base += BK) {
|
||||
const uint buf = (base / BK) & 1u;
|
||||
if (base + BK < p.k)
|
||||
load_tile(base + BK, buf ^ 1u, row, col);
|
||||
coopmat<int8_t, gl_ScopeWorkgroup, BM, BK, gl_MatrixUseA> w;
|
||||
coopmat<int8_t, gl_ScopeWorkgroup, BK, BN, gl_MatrixUseB> a;
|
||||
coopMatLoad(w, weights, buf * BM * STRIDE, STRIDE, gl_CooperativeMatrixLayoutRowMajor);
|
||||
coopMatLoad(a, activations, buf * BN * STRIDE, STRIDE, gl_CooperativeMatrixLayoutColumnMajor);
|
||||
sum = coopMatMulAdd(w, a, sum);
|
||||
barrier();
|
||||
}
|
||||
coopMatPerElementNV(sum, sum, dequantize, row, col);
|
||||
tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> dl = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV);
|
||||
dl = setTensorLayoutDimensionNV(dl, p.rows, p.n);
|
||||
tensorViewNV<2, false, 1, 0> transpose = createTensorViewNV(2, false, 1, 0);
|
||||
coopMatStoreTensorNV(sum, data_d, 0, sliceTensorLayoutNV(dl, row, BN, col, BM), transpose);
|
||||
}
|
||||
@@ -0,0 +1,92 @@
|
||||
#version 460
|
||||
#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require
|
||||
#extension GL_EXT_shader_8bit_storage : require
|
||||
#extension GL_EXT_buffer_reference : require
|
||||
#extension GL_EXT_shader_explicit_arithmetic_types_int32 : require
|
||||
#extension GL_KHR_cooperative_matrix : require
|
||||
#extension GL_NV_cooperative_matrix2 : require
|
||||
#extension GL_KHR_memory_scope_semantics : require
|
||||
layout(local_size_x = 256) in;
|
||||
layout(constant_id = 0) const uint BM = 64;
|
||||
layout(constant_id = 1) const uint BN = 64;
|
||||
layout(constant_id = 2) const uint BK = 64;
|
||||
layout(constant_id = 3) const bool FULL = true;
|
||||
layout(push_constant) uniform parameter {
|
||||
uint n;
|
||||
uint rows;
|
||||
uint k;
|
||||
uint has_bias;
|
||||
uint row_offset;
|
||||
}
|
||||
p;
|
||||
layout(binding = 0, std430) readonly buffer W {
|
||||
uvec4 data_w[];
|
||||
};
|
||||
layout(binding = 1, std430) readonly buffer A {
|
||||
int8_t data_a[];
|
||||
};
|
||||
layout(binding = 2, std430) readonly buffer WS {
|
||||
float data_ws[];
|
||||
};
|
||||
layout(binding = 3, std430) readonly buffer AS {
|
||||
float data_as[];
|
||||
};
|
||||
layout(binding = 4, std430) readonly buffer B {
|
||||
float data_b[];
|
||||
};
|
||||
layout(binding = 5, std430) writeonly buffer D {
|
||||
int32_t data_d[];
|
||||
};
|
||||
layout(buffer_reference, std430, buffer_reference_align = 16) buffer PackedBlock {
|
||||
uvec4 value;
|
||||
};
|
||||
int8_t decode_weight(const in PackedBlock block, const in uint block_coords[2], const in uint coord[2]) {
|
||||
return int8_t(bitfieldExtract(int(block.value[coord[1] / 8]), int((coord[1] % 8) * 4), 4));
|
||||
}
|
||||
int32_t dequantize(const in uint r, const in uint c, const in int32_t v, const in uint row, const in uint col) {
|
||||
if (row + c >= p.rows || col + r >= p.n)
|
||||
return 0;
|
||||
precise float value = float(v) * data_as[row + c] * data_ws[col + r];
|
||||
if (p.has_bias != 0)
|
||||
value += data_b[col + r];
|
||||
return floatBitsToInt(value);
|
||||
}
|
||||
// The host selects complete weight tiles and a 256-aligned input width.
|
||||
void main() {
|
||||
const uint cols = gl_NumWorkGroups.x, rows = gl_NumWorkGroups.y;
|
||||
const uint linear = gl_WorkGroupID.y * cols + gl_WorkGroupID.x;
|
||||
const uint panel = min(32u, cols), first_col = (linear / (panel * rows)) * panel;
|
||||
const uint group_cols = min(panel, cols - first_col), within = linear - first_col * rows;
|
||||
const uint row = p.row_offset + (within / group_cols) * BN, col = (first_col + within % group_cols) * BM;
|
||||
coopmat<int32_t, gl_ScopeWorkgroup, BM, BN, gl_MatrixUseAccumulator> sum = coopmat<int32_t, gl_ScopeWorkgroup, BM, BN, gl_MatrixUseAccumulator>(0);
|
||||
tensorLayoutNV<2> wl = createTensorLayoutNV(2);
|
||||
wl = setTensorLayoutBlockSizeNV(wl, 1, 32);
|
||||
wl = setTensorLayoutDimensionNV(wl, p.n, p.k & ~255u);
|
||||
tensorLayoutNV<2> al = createTensorLayoutNV(2);
|
||||
al = setTensorLayoutDimensionNV(al, p.rows, p.k & ~255u);
|
||||
tensorViewNV<2, false, 1, 0> at = createTensorViewNV(2, false, 1, 0);
|
||||
if (FULL) {
|
||||
for (uint base = 0; base < p.k; base += BK) {
|
||||
coopmat<int8_t, gl_ScopeWorkgroup, BM, BK, gl_MatrixUseA> w;
|
||||
coopmat<int8_t, gl_ScopeWorkgroup, BK, BN, gl_MatrixUseB> a;
|
||||
coopMatLoadTensorNV(w, data_w, 0, sliceTensorLayoutNV(wl, col, BM, base, BK), decode_weight);
|
||||
coopMatLoadTensorNV(a, data_a, 0, sliceTensorLayoutNV(al, row, BN, base, BK), at);
|
||||
sum = coopMatMulAdd(w, a, sum);
|
||||
}
|
||||
} else {
|
||||
tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> tail = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV);
|
||||
tail = setTensorLayoutDimensionNV(tail, p.rows, p.k & ~255u);
|
||||
for (uint base = 0; base < p.k; base += BK) {
|
||||
coopmat<int8_t, gl_ScopeWorkgroup, BM, BK, gl_MatrixUseA> w;
|
||||
coopmat<int8_t, gl_ScopeWorkgroup, BK, BN, gl_MatrixUseB> a;
|
||||
coopMatLoadTensorNV(w, data_w, 0, sliceTensorLayoutNV(wl, col, BM, base, BK), decode_weight);
|
||||
coopMatLoadTensorNV(a, data_a, 0, sliceTensorLayoutNV(tail, row, BN, base, BK), at);
|
||||
sum = coopMatMulAdd(w, a, sum);
|
||||
}
|
||||
}
|
||||
coopMatPerElementNV(sum, sum, dequantize, row, col);
|
||||
tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> dl = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV);
|
||||
dl = setTensorLayoutDimensionNV(dl, p.rows, p.n);
|
||||
tensorViewNV<2, false, 1, 0> transpose = createTensorViewNV(2, false, 1, 0);
|
||||
coopMatStoreTensorNV(sum, data_d, 0, sliceTensorLayoutNV(dl, row, BN, col, BM), transpose);
|
||||
}
|
||||
@@ -0,0 +1,164 @@
|
||||
#version 460
|
||||
|
||||
#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require
|
||||
#extension GL_EXT_shader_explicit_arithmetic_types_int32 : require
|
||||
#extension GL_KHR_cooperative_matrix : require
|
||||
#extension GL_NV_cooperative_matrix2 : require
|
||||
#extension GL_KHR_memory_scope_semantics : require
|
||||
#extension GL_EXT_control_flow_attributes : require
|
||||
|
||||
layout(local_size_x = 256) in;
|
||||
layout(constant_id = 0) const uint BM = 64;
|
||||
layout(constant_id = 1) const uint BN = 64;
|
||||
layout(constant_id = 2) const uint BK = 64;
|
||||
const uint STRIDE = BK / 2 + 4;
|
||||
|
||||
layout(push_constant) uniform parameter {
|
||||
uint n;
|
||||
uint rows;
|
||||
uint k;
|
||||
uint has_bias;
|
||||
}
|
||||
p;
|
||||
|
||||
layout(binding = 0, std430) readonly buffer W {
|
||||
uvec2 data_w[];
|
||||
};
|
||||
layout(binding = 1, std430) readonly buffer A {
|
||||
uvec4 data_a[];
|
||||
};
|
||||
layout(binding = 2, std430) readonly buffer CB {
|
||||
float data_cb[];
|
||||
};
|
||||
layout(binding = 3, std430) readonly buffer SC {
|
||||
float data_sc[];
|
||||
};
|
||||
layout(binding = 4, std430) readonly buffer SR {
|
||||
uint data_sr[];
|
||||
};
|
||||
layout(binding = 5, std430) readonly buffer AS {
|
||||
float data_as[];
|
||||
};
|
||||
layout(binding = 6, std430) readonly buffer B {
|
||||
float data_b[];
|
||||
};
|
||||
layout(binding = 7, std430) writeonly buffer D {
|
||||
int32_t data_d[];
|
||||
};
|
||||
shared uint lut[256];
|
||||
shared int cb[16];
|
||||
shared float cb_scale;
|
||||
shared float relative[256];
|
||||
shared float scales[BK / 16 * BM];
|
||||
|
||||
float fp8_e4m3(uint v) {
|
||||
const uint exponent = (v >> 3) & 15u, mantissa = v & 7u;
|
||||
const float value = exponent == 15u && mantissa == 7u ? uintBitsToFloat(0x7fc00000u) : exponent == 0u ? ldexp(float(mantissa), -9)
|
||||
: ldexp(1.0 + float(mantissa) / 8.0, int(exponent) - 7);
|
||||
return (v & 128u) != 0u ? -value : value;
|
||||
}
|
||||
|
||||
int32_t accumulate_group(const in uint row, const in uint col, const in int32_t value, const in int32_t previous, const in uint group) {
|
||||
return floatBitsToInt(fma(float(value), scales[group * BM + row], intBitsToFloat(previous)));
|
||||
}
|
||||
|
||||
shared uint weights[BM * STRIDE];
|
||||
const uint A_STRIDE = BK / 4 + 4;
|
||||
shared uint activations[BN * A_STRIDE];
|
||||
|
||||
int32_t dequantize(const in uint row, const in uint col, const in int32_t acc, const in uint row_base, const in uint col_base) {
|
||||
const uint r = row_base + col, c = col_base + row;
|
||||
if (r >= p.rows || c >= p.n) {
|
||||
return 0;
|
||||
}
|
||||
precise float value = intBitsToFloat(acc) * data_as[r];
|
||||
if (p.has_bias != 0) {
|
||||
value += data_b[c];
|
||||
}
|
||||
return floatBitsToInt(value);
|
||||
}
|
||||
|
||||
void main() {
|
||||
const uint tid = gl_LocalInvocationIndex;
|
||||
const uint cols = gl_NumWorkGroups.x, rows = gl_NumWorkGroups.y;
|
||||
const uint linear = gl_WorkGroupID.y * cols + gl_WorkGroupID.x;
|
||||
// Reuse weight panels across activation tiles when the output expands.
|
||||
const uint panel = p.n > p.k ? 16u : cols;
|
||||
const uint first_col = (linear / (panel * rows)) * panel;
|
||||
const uint group_cols = min(panel, cols - first_col);
|
||||
const uint within = linear - first_col * rows;
|
||||
const uint col_base = (first_col + within % group_cols) * BM, row_base = (within / group_cols) * BN;
|
||||
tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> d_layout = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV);
|
||||
d_layout = setTensorLayoutDimensionNV(d_layout, p.rows, p.n);
|
||||
tensorViewNV<2, false, 1, 0> transpose = createTensorViewNV(2, false, 1, 0);
|
||||
coopmat<int32_t, gl_ScopeWorkgroup, BM, BN, gl_MatrixUseAccumulator> sum = coopmat<int32_t, gl_ScopeWorkgroup, BM, BN, gl_MatrixUseAccumulator>(0);
|
||||
if (tid == 0) {
|
||||
float amax = 0.0;
|
||||
for (uint i = 0; i < 16; ++i) {
|
||||
amax = max(amax, abs(data_cb[i]));
|
||||
}
|
||||
cb_scale = amax / 127.0;
|
||||
for (uint i = 0; i < 16; ++i) {
|
||||
cb[i] = cb_scale == 0.0 ? 0 : clamp(int(roundEven(data_cb[i] / cb_scale)), -127, 127);
|
||||
}
|
||||
}
|
||||
relative[tid] = fp8_e4m3(tid);
|
||||
barrier();
|
||||
for (uint i = tid; i < 256; i += 256) {
|
||||
lut[i] = uint(cb[i & 15u] & 255) | (uint(cb[i >> 4] & 255) << 8);
|
||||
}
|
||||
barrier();
|
||||
for (uint base = 0; base < p.k; base += BK) {
|
||||
for (uint i = tid; i < BN * BK / 16; i += 256) {
|
||||
const uint r = i / (BK / 16), t = i % (BK / 16);
|
||||
const uvec4 value = row_base + r < p.rows ? data_a[(row_base + r) * (p.k / 16) + base / 16 + t] : uvec4(0);
|
||||
const uint offset = r * A_STRIDE + t * 4;
|
||||
activations[offset] = value.x;
|
||||
activations[offset + 1] = value.y;
|
||||
activations[offset + 2] = value.z;
|
||||
activations[offset + 3] = value.w;
|
||||
}
|
||||
for (uint i = tid; i < BM * BK / 16; i += 256) {
|
||||
const uint c = i / (BK / 16), part = i % (BK / 16);
|
||||
const uvec2 packed = col_base + c < p.n ? data_w[(col_base + c) * (p.k / 16) + base / 16 + part] : uvec2(0);
|
||||
// Complementary zero halves let adjacent groups reuse one activation fragment.
|
||||
const uint odd = part % 2;
|
||||
const uint offset = c * STRIDE + part * 8 + odd * 4;
|
||||
const uint zero_offset = c * STRIDE + part * 8 + (1 - odd) * 4;
|
||||
for (uint j = 0; j < 2; ++j) {
|
||||
const uint v = packed[j];
|
||||
weights[offset + 2 * j] = lut[v & 255u] | (lut[(v >> 8) & 255u] << 16);
|
||||
weights[offset + 2 * j + 1] = lut[(v >> 16) & 255u] | (lut[v >> 24] << 16);
|
||||
weights[zero_offset + 2 * j] = 0u;
|
||||
weights[zero_offset + 2 * j + 1] = 0u;
|
||||
}
|
||||
}
|
||||
for (uint i = tid; i < BM * BK / 16; i += 256) {
|
||||
const uint c = i / (BK / 16), g = i % (BK / 16);
|
||||
float scale = 0.0;
|
||||
if (col_base + c < p.n) {
|
||||
const uint packed = data_sr[(col_base + c) * (p.k / 64) + base / 64 + g / 4];
|
||||
scale = cb_scale * (data_sc[col_base + c] * relative[(packed >> (8 * (g % 4))) & 255u]);
|
||||
}
|
||||
scales[g * BM + c] = scale;
|
||||
}
|
||||
barrier();
|
||||
// The relative scale changes every 16 input values, so each group
|
||||
// needs its own integer MMA and floating-point accumulation. Pad the
|
||||
// weights with zeros to satisfy the device's K=32 matrix granularity.
|
||||
[[unroll]] for (uint g = 0; g < BK / 16; g += 2) {
|
||||
coopmat<int8_t, gl_ScopeWorkgroup, 32, BN, gl_MatrixUseB> a;
|
||||
coopMatLoad(a, activations, g * 4, A_STRIDE, gl_CooperativeMatrixLayoutColumnMajor);
|
||||
[[unroll]] for (uint j = 0; j < 2; ++j) {
|
||||
coopmat<int8_t, gl_ScopeWorkgroup, BM, 32, gl_MatrixUseA> w;
|
||||
coopmat<int32_t, gl_ScopeWorkgroup, BM, BN, gl_MatrixUseAccumulator> part = coopmat<int32_t, gl_ScopeWorkgroup, BM, BN, gl_MatrixUseAccumulator>(0);
|
||||
coopMatLoad(w, weights, (g + j) * 8, STRIDE, gl_CooperativeMatrixLayoutRowMajor);
|
||||
part = coopMatMulAdd(w, a, part);
|
||||
coopMatPerElementNV(sum, part, accumulate_group, sum, g + j);
|
||||
}
|
||||
}
|
||||
barrier();
|
||||
}
|
||||
coopMatPerElementNV(sum, sum, dequantize, row_base, col_base);
|
||||
coopMatStoreTensorNV(sum, data_d, 0, sliceTensorLayoutNV(d_layout, row_base, BN, col_base, BM), transpose);
|
||||
}
|
||||
Reference in new issue
Block a user