Skip to content

vulkan: use spec constant for matrix matrix multiplication A-type #25773

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
0cc4m wants to merge 9 commits into
base: master
Choose a base branch
from
Open

vulkan: use spec constant for matrix matrix multiplication A-type #25773

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
c25de61
vulkan: use spec constant for mul mat type_a
0cc4m Jul 7, 2026
b1e0349
remove LUT quants from unified shader
0cc4m Aug 26, 2026
f58bd0a
clean up
0cc4m Aug 27, 2026
6537ccf
restore coopmat2 q4_k/q5_k optimization
0cc4m Aug 28, 2026
8427716
split out q4_k/q5_k cm2 shader to fix Ampere regression
0cc4m Aug 28, 2026
36be26c
revert iq shmem table renames
0cc4m Aug 31, 2026
c14706c
simplify cm2 code with single uint8_t buffer
0cc4m Aug 31, 2026
1ac1ad5
fix fp4 extension use switch being overwritten by generic shader
0cc4m Aug 31, 2026
b7b2b39
clean up
0cc4m Sep 3, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1,595 changes: 648 additions & 947 deletions ggml/src/ggml-vulkan/ggml-vulkan.cpp
View file
Open in desktop

Large diffs are not rendered by default.

8 changes: 4 additions & 4 deletions ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs_cm2.glsl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1041,7 +1041,7 @@ float16_t dequantFuncIQ2_S(const in decodeBufIQ2_S bl, const in uint blockCoords
const uint scale = (bl.block.scales[ib32] >> ((idx & 0x10) >> 2)) & 0xf;
const uint qs = bl.block.qs[ib8];
const uint qh = bl.block.qh[ib32];
const uint sign = bl.block.qs[QUANT_K / 8 + ib8] >> (idx & 0x6);
const uint sign = bl.block.qs[QUANT_K_IQ2_S / 8 + ib8] >> (idx & 0x6);

const float d = float(bl.block.d);
const float db = d * 0.25 * (0.5 + scale);
Expand All @@ -1063,7 +1063,7 @@ f16vec4 dequantFuncIQ2_S_v(const in decodeBufIQ2_S bl, const in uint blockCoords
const uint scale = (bl.block.scales[ib32] >> ((idx & 0x10) >> 2)) & 0xf;
const uint qs = bl.block.qs[ib8];
const uint qh = bl.block.qh[ib32];
const uint sb = uint(bl.block.qs[QUANT_K / 8 + ib8]) >> (idx & 0x6u);
const uint sb = uint(bl.block.qs[QUANT_K_IQ2_S / 8 + ib8]) >> (idx & 0x6u);

const float d = float(bl.block.d);
const float db = d * 0.25 * (0.5 + scale);
Expand Down Expand Up @@ -1094,7 +1094,7 @@ float16_t dequantFuncIQ3_XXS(const in decodeBufIQ3_XXS bl, const in uint blockCo
uint idx = coordInBlock[1];

const uint iqs = (idx & 0xFC) >> 2; // 0..63
const uint is = QUANT_K / 4 + ((idx & 0xE0) >> 3);// 8 values
const uint is = QUANT_K_IQ3_XXS / 4 + ((idx & 0xE0) >> 3);// 8 values

const float d = float(bl.block.d);
const uint qs = bl.block.qs[iqs];
Expand All @@ -1117,7 +1117,7 @@ f16vec4 dequantFuncIQ3_XXS_v(const in decodeBufIQ3_XXS bl, const in uint blockCo
const uint idx = coordInBlock[1];

const uint iqs = idx >> 2;
const uint is = QUANT_K / 4 + ((idx & 0xE0) >> 3);
const uint is = QUANT_K_IQ3_XXS / 4 + ((idx & 0xE0) >> 3);

const float d = float(bl.block.d);
const uint qs = bl.block.qs[iqs];
Expand Down
42 changes: 16 additions & 26 deletions ggml/src/ggml-vulkan/vulkan-shaders/fa_types.glsl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,32 +1,22 @@
#if !defined(GGML_FA_TYPES_COMP)
#define GGML_FA_TYPES_COMP

// FaTypeK / FaTypeV spec constant values. These mirror enum ggml_type so the
// host can pass the type directly. Keep in sync with ggml.h.
#define FA_TYPE_F32 0u
#define FA_TYPE_F16 1u
#define FA_TYPE_Q4_0 2u
#define FA_TYPE_Q4_1 3u
#define FA_TYPE_Q5_0 6u
#define FA_TYPE_Q5_1 7u
#define FA_TYPE_Q8_0 8u
#define FA_TYPE_IQ4_NL 20u
#define FA_TYPE_BF16 30u
#include "ggml_type_ids.glsl"

// Number of matrix elements per buffer block, derived from the K/V type spec
// constant. F32 is treated as a vec4 "block" of 4 floats. F16 uses block size 1
// and bypasses the dequant path entirely. Quants follow their ggml block sizes.
uint fa_block_elems(uint ty) {
switch (ty) {
case FA_TYPE_F32: return 4u;
case FA_TYPE_F16: return 1u;
case FA_TYPE_Q4_0: return uint(QUANT_K_Q4_0);
case FA_TYPE_Q4_1: return uint(QUANT_K_Q4_1);
case FA_TYPE_Q5_0: return uint(QUANT_K_Q5_0);
case FA_TYPE_Q5_1: return uint(QUANT_K_Q5_1);
case FA_TYPE_Q8_0: return uint(QUANT_K_Q8_0);
case FA_TYPE_IQ4_NL: return uint(QUANT_K_IQ4_NL);
case FA_TYPE_BF16: return 1u;
case GGML_TYPE_F32: return 4u;
case GGML_TYPE_F16: return 1u;
case GGML_TYPE_Q4_0: return uint(QUANT_K_Q4_0);
case GGML_TYPE_Q4_1: return uint(QUANT_K_Q4_1);
case GGML_TYPE_Q5_0: return uint(QUANT_K_Q5_0);
case GGML_TYPE_Q5_1: return uint(QUANT_K_Q5_1);
case GGML_TYPE_Q8_0: return uint(QUANT_K_Q8_0);
case GGML_TYPE_IQ4_NL: return uint(QUANT_K_IQ4_NL);
case GGML_TYPE_BF16: return 1u;
default: return 1u;
}
}
Expand All @@ -36,18 +26,18 @@ uint fa_block_elems(uint ty) {
// of int32s per 32-element block on the MMQ K path: ints_per_block == 8 / R.
uint fa_quant_r_mmq(uint ty) {
switch (ty) {
case FA_TYPE_Q4_0: return uint(QUANT_R_Q4_0);
case FA_TYPE_Q4_1: return uint(QUANT_R_Q4_1);
case FA_TYPE_Q5_0: return uint(QUANT_R_Q5_0);
case FA_TYPE_Q5_1: return uint(QUANT_R_Q5_1);
case FA_TYPE_Q8_0: return uint(QUANT_R_Q8_0);
case GGML_TYPE_Q4_0: return uint(QUANT_R_Q4_0);
case GGML_TYPE_Q4_1: return uint(QUANT_R_Q4_1);
case GGML_TYPE_Q5_0: return uint(QUANT_R_Q5_0);
case GGML_TYPE_Q5_1: return uint(QUANT_R_Q5_1);
case GGML_TYPE_Q8_0: return uint(QUANT_R_Q8_0);
default: return 1u;
}
}

bool fa_type_needs_shmem(uint ty) {
switch (ty) {
case FA_TYPE_IQ4_NL: return true;
case GGML_TYPE_IQ4_NL: return true;
default: return false;
}
}
Expand Down
6 changes: 3 additions & 3 deletions ggml/src/ggml-vulkan/vulkan-shaders/flash_attn.comp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -134,7 +134,7 @@ void main() {

// Q8_0 K only needs (qd, _); the asymmetric Q4_*/Q5_* family also stores
// the row-sum scaled by qd, used in k_dot_correction.
if (FaTypeK == FA_TYPE_Q8_0) {
if (FaTypeK == GGML_TYPE_Q8_0) {
if (buf_iqs == 0) {
Qf[buf_ib].ds = FLOAT_TYPEV2(qd, 0.0f);
}
Expand Down Expand Up @@ -367,15 +367,15 @@ void main() {
// Q4_*/Q5_* take the block-8 fast path when one step covers a full
// block; Q8_0 always goes through the per-int get_k_qs* helpers
// (its qs is byte-packed, not nibble-packed).
const bool block8_fast = (d_per_step == 8) && (FaTypeK != FA_TYPE_Q8_0);
const bool block8_fast = (d_per_step == 8) && (FaTypeK != GGML_TYPE_Q8_0);

if (SHMEM_STAGING != 0) {
const uint k_block_idx = (d_tid * (HSK_per_thread / 4) + d_block) / 8;
const uint buf_ib = (c * cols_per_iter + col_tid) * qf_stride + k_block_idx;
k_dm = ACC_TYPEV2(kblocksh[buf_ib].dm);

if (block8_fast) {
const bool has_qh = (FaTypeK == FA_TYPE_Q5_0) || (FaTypeK == FA_TYPE_Q5_1);
const bool has_qh = (FaTypeK == GGML_TYPE_Q5_0) || (FaTypeK == GGML_TYPE_Q5_1);
[[unroll]] for (uint32_t d = 0; d < 4; d++) {
uint vui = kblocksh[buf_ib].qs[d];
k_quants[d ] = int32_t( vui & 0x0F0F0F0F);
Expand Down
4 changes: 2 additions & 2 deletions ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -105,8 +105,8 @@ layout (binding = 6) readonly buffer MO {uint32_t data_mask_opt[];};
#define BLOCK_SIZE_V fa_block_elems(FaTypeV)
// F16 reads f16 elements directly from the binding; everything else routes
// through dequantize4 / the MMQ helpers to unpack from the packed block layout.
#define USE_DECODE_K (FaTypeK != FA_TYPE_F16)
#define USE_DECODE_V (FaTypeV != FA_TYPE_F16)
#define USE_DECODE_K (FaTypeK != GGML_TYPE_F16)
#define USE_DECODE_V (FaTypeV != GGML_TYPE_F16)

#define CEIL_DIV(a, b) (((a) + (b) - 1) / (b))

Expand Down
58 changes: 30 additions & 28 deletions ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,53 +40,55 @@ layout(buffer_reference, std430, buffer_reference_align = 1) buffer decodeBufFA_
#if !defined(BFLOAT16)
float16_t faDecodeK(const decodeBufFA_K bl_in, const uint blockCoords[2], const uint coordInBlock[2]) {
switch (FaTypeK) {
case FA_TYPE_F32: return dequantFuncF32 (decodeBufF32 (bl_in), blockCoords, coordInBlock);
case FA_TYPE_Q4_0: return dequantFuncQ4_0(decodeBufQ4_0(bl_in), blockCoords, coordInBlock);
case FA_TYPE_Q4_1: return dequantFuncQ4_1(decodeBufQ4_1(bl_in), blockCoords, coordInBlock);
case FA_TYPE_Q5_0: return dequantFuncQ5_0(decodeBufQ5_0(bl_in), blockCoords, coordInBlock);
case FA_TYPE_Q5_1: return dequantFuncQ5_1(decodeBufQ5_1(bl_in), blockCoords, coordInBlock);
case FA_TYPE_Q8_0: return dequantFuncQ8_0(decodeBufQ8_0(bl_in), blockCoords, coordInBlock);
case FA_TYPE_IQ4_NL: return dequantFuncIQ4_NL(decodeBufIQ4_NL(bl_in), blockCoords, coordInBlock);
case GGML_TYPE_F32: return dequantFuncF32 (decodeBufF32 (bl_in), blockCoords, coordInBlock);
case GGML_TYPE_Q4_0: return dequantFuncQ4_0(decodeBufQ4_0(bl_in), blockCoords, coordInBlock);
case GGML_TYPE_Q4_1: return dequantFuncQ4_1(decodeBufQ4_1(bl_in), blockCoords, coordInBlock);
case GGML_TYPE_Q5_0: return dequantFuncQ5_0(decodeBufQ5_0(bl_in), blockCoords, coordInBlock);
case GGML_TYPE_Q5_1: return dequantFuncQ5_1(decodeBufQ5_1(bl_in), blockCoords, coordInBlock);
case GGML_TYPE_Q8_0: return dequantFuncQ8_0(decodeBufQ8_0(bl_in), blockCoords, coordInBlock);
case GGML_TYPE_IQ4_NL: return dequantFuncIQ4_NL(decodeBufIQ4_NL(bl_in), blockCoords, coordInBlock);
case GGML_TYPE_Q1_0: return dequantFuncQ1_0(decodeBufQ1_0(bl_in), blockCoords, coordInBlock);
default: return float16_t(0);
}
}

float16_t faDecodeV(const decodeBufFA_V bl_in, const uint blockCoords[2], const uint coordInBlock[2]) {
switch (FaTypeV) {
case FA_TYPE_F32: return dequantFuncF32 (decodeBufF32 (bl_in), blockCoords, coordInBlock);
case FA_TYPE_Q4_0: return dequantFuncQ4_0(decodeBufQ4_0(bl_in), blockCoords, coordInBlock);
case FA_TYPE_Q4_1: return dequantFuncQ4_1(decodeBufQ4_1(bl_in), blockCoords, coordInBlock);
case FA_TYPE_Q5_0: return dequantFuncQ5_0(decodeBufQ5_0(bl_in), blockCoords, coordInBlock);
case FA_TYPE_Q5_1: return dequantFuncQ5_1(decodeBufQ5_1(bl_in), blockCoords, coordInBlock);
case FA_TYPE_Q8_0: return dequantFuncQ8_0(decodeBufQ8_0(bl_in), blockCoords, coordInBlock);
case FA_TYPE_IQ4_NL: return dequantFuncIQ4_NL(decodeBufIQ4_NL(bl_in), blockCoords, coordInBlock);
case GGML_TYPE_F32: return dequantFuncF32 (decodeBufF32 (bl_in), blockCoords, coordInBlock);
case GGML_TYPE_Q4_0: return dequantFuncQ4_0(decodeBufQ4_0(bl_in), blockCoords, coordInBlock);
case GGML_TYPE_Q4_1: return dequantFuncQ4_1(decodeBufQ4_1(bl_in), blockCoords, coordInBlock);
case GGML_TYPE_Q5_0: return dequantFuncQ5_0(decodeBufQ5_0(bl_in), blockCoords, coordInBlock);
case GGML_TYPE_Q5_1: return dequantFuncQ5_1(decodeBufQ5_1(bl_in), blockCoords, coordInBlock);
case GGML_TYPE_Q8_0: return dequantFuncQ8_0(decodeBufQ8_0(bl_in), blockCoords, coordInBlock);
case GGML_TYPE_IQ4_NL: return dequantFuncIQ4_NL(decodeBufIQ4_NL(bl_in), blockCoords, coordInBlock);
case GGML_TYPE_Q1_0: return dequantFuncQ1_0(decodeBufQ1_0(bl_in), blockCoords, coordInBlock);
default: return float16_t(0);
}
}

// V=4 vector decode for K/V; dispatches to per-format _v decoders.
f16vec4 faDecodeKVector(const decodeBufFA_K bl_in, const uint blockCoords[2], const uint coordInBlock[2]) {
switch (FaTypeK) {
case FA_TYPE_F32: return f16vec4(decodeBufF32(bl_in).block);
case FA_TYPE_Q4_0: return dequantFuncQ4_0_v(decodeBufQ4_0(bl_in), blockCoords, coordInBlock);
case FA_TYPE_Q4_1: return dequantFuncQ4_1_v(decodeBufQ4_1(bl_in), blockCoords, coordInBlock);
case FA_TYPE_Q5_0: return dequantFuncQ5_0_v(decodeBufQ5_0(bl_in), blockCoords, coordInBlock);
case FA_TYPE_Q5_1: return dequantFuncQ5_1_v(decodeBufQ5_1(bl_in), blockCoords, coordInBlock);
case FA_TYPE_Q8_0: return dequantFuncQ8_0_v(decodeBufQ8_0(bl_in), blockCoords, coordInBlock);
case FA_TYPE_IQ4_NL: return dequantFuncIQ4_NL_v(decodeBufIQ4_NL(bl_in), blockCoords, coordInBlock);
case GGML_TYPE_F32: return f16vec4(decodeBufF32(bl_in).block);
case GGML_TYPE_Q4_0: return dequantFuncQ4_0_v(decodeBufQ4_0(bl_in), blockCoords, coordInBlock);
case GGML_TYPE_Q4_1: return dequantFuncQ4_1_v(decodeBufQ4_1(bl_in), blockCoords, coordInBlock);
case GGML_TYPE_Q5_0: return dequantFuncQ5_0_v(decodeBufQ5_0(bl_in), blockCoords, coordInBlock);
case GGML_TYPE_Q5_1: return dequantFuncQ5_1_v(decodeBufQ5_1(bl_in), blockCoords, coordInBlock);
case GGML_TYPE_Q8_0: return dequantFuncQ8_0_v(decodeBufQ8_0(bl_in), blockCoords, coordInBlock);
case GGML_TYPE_IQ4_NL: return dequantFuncIQ4_NL_v(decodeBufIQ4_NL(bl_in), blockCoords, coordInBlock);
default: return f16vec4(0);
}
}

f16vec4 faDecodeVVector(const decodeBufFA_V bl_in, const uint blockCoords[2], const uint coordInBlock[2]) {
switch (FaTypeV) {
case FA_TYPE_F32: return f16vec4(decodeBufF32(bl_in).block);
case FA_TYPE_Q4_0: return dequantFuncQ4_0_v(decodeBufQ4_0(bl_in), blockCoords, coordInBlock);
case FA_TYPE_Q4_1: return dequantFuncQ4_1_v(decodeBufQ4_1(bl_in), blockCoords, coordInBlock);
case FA_TYPE_Q5_0: return dequantFuncQ5_0_v(decodeBufQ5_0(bl_in), blockCoords, coordInBlock);
case FA_TYPE_Q5_1: return dequantFuncQ5_1_v(decodeBufQ5_1(bl_in), blockCoords, coordInBlock);
case FA_TYPE_Q8_0: return dequantFuncQ8_0_v(decodeBufQ8_0(bl_in), blockCoords, coordInBlock);
case FA_TYPE_IQ4_NL: return dequantFuncIQ4_NL_v(decodeBufIQ4_NL(bl_in), blockCoords, coordInBlock);
case GGML_TYPE_F32: return f16vec4(decodeBufF32(bl_in).block);
case GGML_TYPE_Q4_0: return dequantFuncQ4_0_v(decodeBufQ4_0(bl_in), blockCoords, coordInBlock);
case GGML_TYPE_Q4_1: return dequantFuncQ4_1_v(decodeBufQ4_1(bl_in), blockCoords, coordInBlock);
case GGML_TYPE_Q5_0: return dequantFuncQ5_0_v(decodeBufQ5_0(bl_in), blockCoords, coordInBlock);
case GGML_TYPE_Q5_1: return dequantFuncQ5_1_v(decodeBufQ5_1(bl_in), blockCoords, coordInBlock);
case GGML_TYPE_Q8_0: return dequantFuncQ8_0_v(decodeBufQ8_0(bl_in), blockCoords, coordInBlock);
case GGML_TYPE_IQ4_NL: return dequantFuncIQ4_NL_v(decodeBufIQ4_NL(bl_in), blockCoords, coordInBlock);
default: return f16vec4(0);
}
}
Expand Down
32 changes: 16 additions & 16 deletions ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_dequant.glsl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -121,25 +121,25 @@ layout (binding = 1) readonly buffer K_PACKED_Q5_1_P32 { block_q5_1_packed32 dat
FLOAT_TYPEV4 dequantize4(uint ib, uint iqs, uint a_offset, uint binding_idx) {
if (binding_idx == BINDING_IDX_K) {
switch (FaTypeK) {
case FA_TYPE_F32: FA_DEQUANT4_F32 (k_packed_f32)
case FA_TYPE_Q4_0: FA_DEQUANT4_Q4_0(k_packed_q4_0)
case FA_TYPE_Q4_1: FA_DEQUANT4_Q4_1(k_packed_q4_1)
case FA_TYPE_Q5_0: FA_DEQUANT4_Q5_0(k_packed_q5_0)
case FA_TYPE_Q5_1: FA_DEQUANT4_Q5_1(k_packed_q5_1)
case FA_TYPE_Q8_0: FA_DEQUANT4_Q8_0(k_packed_q8_0)
case FA_TYPE_IQ4_NL: FA_DEQUANT4_IQ4_NL(k_packed_iq4_nl)
case FA_TYPE_BF16: FA_DEQUANT4_BF16(k_packed_bf16)
case GGML_TYPE_F32: FA_DEQUANT4_F32 (k_packed_f32)
case GGML_TYPE_Q4_0: FA_DEQUANT4_Q4_0(k_packed_q4_0)
case GGML_TYPE_Q4_1: FA_DEQUANT4_Q4_1(k_packed_q4_1)
case GGML_TYPE_Q5_0: FA_DEQUANT4_Q5_0(k_packed_q5_0)
case GGML_TYPE_Q5_1: FA_DEQUANT4_Q5_1(k_packed_q5_1)
case GGML_TYPE_Q8_0: FA_DEQUANT4_Q8_0(k_packed_q8_0)
case GGML_TYPE_IQ4_NL: FA_DEQUANT4_IQ4_NL(k_packed_iq4_nl)
case GGML_TYPE_BF16: FA_DEQUANT4_BF16(k_packed_bf16)
}
} else {
switch (FaTypeV) {
case FA_TYPE_F32: FA_DEQUANT4_F32 (v_packed_f32)
case FA_TYPE_Q4_0: FA_DEQUANT4_Q4_0(v_packed_q4_0)
case FA_TYPE_Q4_1: FA_DEQUANT4_Q4_1(v_packed_q4_1)
case FA_TYPE_Q5_0: FA_DEQUANT4_Q5_0(v_packed_q5_0)
case FA_TYPE_Q5_1: FA_DEQUANT4_Q5_1(v_packed_q5_1)
case FA_TYPE_Q8_0: FA_DEQUANT4_Q8_0(v_packed_q8_0)
case FA_TYPE_IQ4_NL: FA_DEQUANT4_IQ4_NL(v_packed_iq4_nl)
case FA_TYPE_BF16: FA_DEQUANT4_BF16(v_packed_bf16)
case GGML_TYPE_F32: FA_DEQUANT4_F32 (v_packed_f32)
case GGML_TYPE_Q4_0: FA_DEQUANT4_Q4_0(v_packed_q4_0)
case GGML_TYPE_Q4_1: FA_DEQUANT4_Q4_1(v_packed_q4_1)
case GGML_TYPE_Q5_0: FA_DEQUANT4_Q5_0(v_packed_q5_0)
case GGML_TYPE_Q5_1: FA_DEQUANT4_Q5_1(v_packed_q5_1)
case GGML_TYPE_Q8_0: FA_DEQUANT4_Q8_0(v_packed_q8_0)
case GGML_TYPE_IQ4_NL: FA_DEQUANT4_IQ4_NL(v_packed_iq4_nl)
case GGML_TYPE_BF16: FA_DEQUANT4_BF16(v_packed_bf16)
}
}
return FLOAT_TYPEV4(0);
Expand Down
Loading
Loading