Lumpiasty/llama.cpp
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Feature: Added LIGHTNING_INDEXER support for Deepseek V4 ops on Vulkan Backend (#27453)

Browse Source 
* vulkan: add LIGHTNING_INDEXER op

* vulkan: updated lightning_indexer.comp and ggml-vulkan.cpp with 128-lane dot-product reduction moved from a shared-memory tree to subgroupAdd.

* vulkan: cleanup; Skip bounds checks

* vulkan: cleanup FA_K_ONLY

* Revert "vulkan: cleanup FA_K_ONLY"

This reverts commit fdcbdd91511945d6878d9070b5447e4a34dce010.

* vulkan: restore interleaved K/V buffer ordering

* vulkan: Remove FA_K_ONLY

* vulkan: Revert flash_attn_dequant

* vulkan: Revert tests in backend-ops.cpp
This commit is contained in:
Shobhit
2026-08-27 15:34:42 +02:00
committed by GitHub
parent 1a946ec745
commit cb300598d5
5 changed files with 365 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-vulkan/vulkan-shaders/fa_types.glsl
+55
View File
@@ -0,0 +1,55 @@
#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
// 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;
default: return 1u;
}
}
// QUANT_R_MMQ for FA-eligible K types. Q4_*/Q5_* store two nibbles per byte
// (R==2); Q8_0 stores one byte per element (R==1). Used to derive the number
// 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);
default: return 1u;
}
}
bool fa_type_needs_shmem(uint ty) {
switch (ty) {
case FA_TYPE_IQ4_NL: return true;
default: return false;
}
}
#endif // !defined(GGML_FA_TYPES_COMP)
ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl
+1 -50
View File
@@ -88,17 +88,7 @@ layout (binding = 6) readonly buffer MO {uint32_t data_mask_opt[];};
#define BINDING_IDX_K 0
#define BINDING_IDX_V 1
// 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 "fa_types.glsl"
#if defined(BFLOAT16)
#define O_TYPE float
@@ -108,45 +98,6 @@ layout (binding = 6) readonly buffer MO {uint32_t data_mask_opt[];};
#define O_TYPEV4 FLOAT_TYPEV4
#endif
// 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;
default: return 1u;
}
}
// QUANT_R_MMQ for FA-eligible K types. Q4_*/Q5_* store two nibbles per byte
// (R==2); Q8_0 stores one byte per element (R==1). Used to derive the number
// 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);
default: return 1u;
}
}
bool fa_type_needs_shmem(uint ty) {
switch (ty) {
case FA_TYPE_IQ4_NL: return true;
default: return false;
}
}
// These can't be `const` globals because GLSL forbids function calls in global
// const initializers, even when the spec constants would let the driver fold
// them. Macros expand at the use site and fold after specialization.
ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer.comp
+151
View File
@@ -0,0 +1,151 @@
#version 450
#extension GL_EXT_control_flow_attributes : require
#extension GL_EXT_shader_16bit_storage : require
#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
#extension GL_KHR_shader_subgroup_basic : enable
#if USE_SUBGROUP_ADD
#extension GL_KHR_shader_subgroup_arithmetic : enable
#endif
#define BINDING_IDX_K 0u
#include "types.glsl"
#include "fa_types.glsl"
#define FaTypeV FA_TYPE_F32
layout(constant_id = 0) const uint FaTypeK = FA_TYPE_F32;
layout(constant_id = 1) const uint FaBlockBytesK = 4;
layout(constant_id = 2) const uint SUBGROUP_SIZE = 32;
#include "flash_attn_dequant.glsl"
// one workgroup computes one output element, one invocation per head element
#define HEAD_SIZE 128
layout(local_size_x = HEAD_SIZE, local_size_y = 1, local_size_z = 1) in;
layout(binding = 0) readonly buffer QBuf { float q[]; };
layout(binding = 1) readonly buffer KBufF16 { float16_t k_f16[]; };
layout(binding = 1) readonly buffer KBufF32 { float k_f32[]; };
layout(binding = 1) readonly buffer KBufBF16 { uint16_t k_bf16[]; };
layout(binding = 2) readonly buffer WBuf { float weights[]; };
layout(binding = 3) readonly buffer MBuf { float16_t mask[]; };
layout(binding = 4) writeonly buffer DstBuf { float dst[]; };
layout(push_constant) uniform PushConstants {
uint n_kv;
uint n_heads;
uint n_tokens;
uint n_streams;
uint n_masks;
uint dispatch_x;
uint q_nb1;
uint q_nb2;
uint q_nb3;
uint k_nb2;
uint k_nb3;
uint w_nb1;
uint w_nb3;
uint m_nb1;
uint m_nb3;
uint d_nb1;
uint d_nb3;
};
shared float k_row[HEAD_SIZE];
#if USE_SUBGROUP_ADD
shared float sg_partials[HEAD_SIZE / SUBGROUP_SIZE];
#else
shared float partials[HEAD_SIZE];
#endif
void main() {
const uint tid = gl_LocalInvocationID.x;
const uint output_idx = gl_WorkGroupID.y * dispatch_x + gl_WorkGroupID.x;
const uint n_outputs = n_kv * n_tokens * n_streams;
if (fa_type_needs_shmem(FaTypeK)) {
init_iq_shmem(gl_WorkGroupSize);
}
if (output_idx >= n_outputs) {
return;
}
const uint ik = output_idx % n_kv;
const uint ts = output_idx / n_kv;
const uint t = ts % n_tokens;
const uint s = ts / n_tokens;
const uint k_offset = ik * k_nb2 + s * k_nb3;
// k strides come in as bytes, so scale them down to the view being indexed
const uint k_block_elems = fa_block_elems(FaTypeK);
const uint k_elem_bytes = FaBlockBytesK / k_block_elems;
if (FaTypeK == FA_TYPE_F16) {
k_row[tid] = float(k_f16[k_offset / k_elem_bytes + tid]);
} else if (FaTypeK == FA_TYPE_F32) {
k_row[tid] = k_f32[k_offset / k_elem_bytes + tid];
} else if (FaTypeK == FA_TYPE_BF16) {
k_row[tid] = bf16_to_fp32(uint(k_bf16[k_offset / k_elem_bytes + tid]));
} else if (4 * tid < HEAD_SIZE) {
const uint coord = 4 * tid;
const uint ib = coord / k_block_elems;
const uint iqs = coord % k_block_elems;
const vec4 values = dequantize4(ib, iqs, k_offset / FaBlockBytesK, BINDING_IDX_K);
k_row[coord + 0] = values.x;
k_row[coord + 1] = values.y;
k_row[coord + 2] = values.z;
k_row[coord + 3] = values.w;
}
barrier();
const float k_val = k_row[tid];
float score = 0.0;
for (uint h = 0; h < n_heads; ++h) {
const float prod = q[h * q_nb1 + t * q_nb2 + s * q_nb3 + tid] * k_val;
#if USE_SUBGROUP_ADD
const float sg_sum = subgroupAdd(prod);
if (gl_SubgroupInvocationID == 0) {
sg_partials[gl_SubgroupID] = sg_sum;
}
barrier();
if (tid == 0) {
float sum = 0.0;
[[unroll]] for (uint i = 0; i < HEAD_SIZE / SUBGROUP_SIZE; ++i) {
sum += sg_partials[i];
}
score += max(sum, 0.0) * weights[h + t * w_nb1 + s * w_nb3];
}
// the reads above must complete before the next iteration overwrites sg_partials
barrier();
#else
partials[tid] = prod;
barrier();
[[unroll]] for (uint stride = HEAD_SIZE / 2; stride > 0; stride >>= 1) {
if (tid < stride) {
partials[tid] += partials[tid + stride];
}
barrier();
}
if (tid == 0) {
score += max(partials[0], 0.0) * weights[h + t * w_nb1 + s * w_nb3];
}
// the read of partials[0] above must complete before the next iteration
// overwrites partials[tid]
barrier();
#endif
}
if (tid == 0) {
const uint mask_offset = ik + t * m_nb1 + (s % n_masks) * m_nb3;
dst[ik + t * d_nb1 + s * d_nb3] = score + float(mask[mask_offset]);
}
}
ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp
+6
View File
@@ -1069,6 +1069,12 @@ void process_shaders() {
string_to_spv("gated_linear_attn_f32", "gla.comp", merge_maps(base_dict, {{"A_TYPE", "float"}}));
// Compile IQ4_NL support in so its shared LUT is available when K uses it.
// K quant type is selected at runtime via the FaTypeK spec constant.
std::map<std::string, std::string> li_dict = {{"FLOAT_TYPE", "float"}, {"FLOAT_TYPEV4", "vec4"}, {"DATA_A_IQ4_NL", "1"}};
string_to_spv("lightning_indexer_f32", "lightning_indexer.comp", li_dict);
string_to_spv("lightning_indexer_subgroup_f32", "lightning_indexer.comp", merge_maps(li_dict, {{"USE_SUBGROUP_ADD", "1"}}));
string_to_spv("rwkv_wkv7_f32", "wkv7.comp", merge_maps(base_dict, {{"A_TYPE", "float"}}));
string_to_spv("gated_delta_net_f32", "gated_delta_net.comp", merge_maps(base_dict, {{"FLOAT_TYPE", "float"}, {"USE_SUBGROUP_ADD", "1"}, {"USE_SUBGROUP_CLUSTERED", "1"}}));