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

* 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 fdcbdd9151.

* 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
+152 -4
View File
@@ -767,6 +767,21 @@ static constexpr std::initializer_list<std::array<int, 3>> rms_norm_mul_rope_vie
{ 4, 0, 3 }, // set_rows->src[0] == view
};
static constexpr std::array<ggml_type, 9> lightning_indexer_k_types = {
GGML_TYPE_F32,
GGML_TYPE_F16,
GGML_TYPE_BF16,
GGML_TYPE_Q8_0,
GGML_TYPE_Q5_1,
GGML_TYPE_Q5_0,
GGML_TYPE_Q4_1,
GGML_TYPE_Q4_0,
GGML_TYPE_IQ4_NL,
};
static bool ggml_vk_lightning_indexer_k_type_supported(ggml_type type) {
return std::find(lightning_indexer_k_types.begin(), lightning_indexer_k_types.end(), type) != lightning_indexer_k_types.end();
}
struct vk_device_struct {
std::recursive_mutex mutex;
@@ -1068,6 +1083,7 @@ struct vk_device_struct {
vk_pipeline pipeline_rwkv_wkv6_f32;
vk_pipeline pipeline_rwkv_wkv7_f32;
vk_pipeline pipeline_gated_linear_attn_f32;
vk_pipeline pipeline_lightning_indexer_f32[GGML_TYPE_COUNT];
// [size_idx][kda] where size_idx: 0=d16, 1=d32, 2=d64, 3=d128
vk_pipeline pipeline_gated_delta_net[4][2];
vk_pipeline pipeline_ssm_scan_f32_d128;
@@ -1848,6 +1864,26 @@ struct vk_op_gated_linear_attn_push_constants {
uint32_t H;
float scale;
};
struct vk_op_lightning_indexer_push_constants {
uint32_t n_kv;
uint32_t n_heads;
uint32_t n_tokens;
uint32_t n_streams;
uint32_t n_masks;
uint32_t dispatch_x;
uint32_t q_nb1;
uint32_t q_nb2;
uint32_t q_nb3;
uint32_t k_nb2;
uint32_t k_nb3;
uint32_t w_nb1;
uint32_t w_nb3;
uint32_t m_nb1;
uint32_t m_nb3;
uint32_t d_nb1;
uint32_t d_nb3;
};
static_assert(sizeof(vk_op_lightning_indexer_push_constants) <= 128);
struct vk_op_gated_delta_net_push_constants {
uint32_t H;
uint32_t n_tokens;
@@ -3904,11 +3940,16 @@ static vk_fa_pipeline_state get_fa_pipeline_state(const vk_device& device, const
return vk_fa_pipeline_state{hsk, hsv, params.block_rows, params.block_cols, params.d_split, params.row_split, params.shmem_staging, params.path, params.workgroup_size, subgroup_size, aligned, f32acc, flags, params.limit_occupancy_shmem, k_type, v_type};
}
// Bytes per buffer block for the FaBlockBytesK/V spec constants. F32 is fed as
// a vec4 "block" of 4 floats, everything else uses its ggml block size.
static uint32_t fa_block_bytes(ggml_type t) {
if (t == GGML_TYPE_F32) {
return 16u;
}
return (uint32_t) ggml_type_size(t);
}
static std::vector<uint32_t> get_fa_spec_constants(const vk_fa_pipeline_state& state) {
const auto fa_block_bytes = [](ggml_type t) -> uint32_t {
if (t == GGML_TYPE_F32) return 16u;
return (uint32_t) ggml_type_size(t);
};
return {
/* 0 WorkGroupSize */ state.workgroup_size,
/* 1 Br */ state.Br,
@@ -5847,6 +5888,17 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
ggml_vk_create_pipeline(device, device->pipeline_gated_linear_attn_f32, "gated_linear_attn_f32", gated_linear_attn_f32_len, gated_linear_attn_f32_data, "main", 6, sizeof(vk_op_gated_linear_attn_push_constants), {1, 1, 1}, {}, 1);
{
const bool li_subgroup = device->subgroup_arithmetic && device->subgroup_require_full_support;
const size_t li_len = li_subgroup ? lightning_indexer_subgroup_f32_len : lightning_indexer_f32_len;
const void * li_data = li_subgroup ? (const void *)lightning_indexer_subgroup_f32_data : (const void *)lightning_indexer_f32_data;
for (ggml_type k_type : lightning_indexer_k_types) {
const std::string name = "lightning_indexer_" + std::string(ggml_type_name(k_type)) + "_k_f32";
ggml_vk_create_pipeline(device, device->pipeline_lightning_indexer_f32[k_type], name.c_str(), li_len, li_data, "main", 5, sizeof(vk_op_lightning_indexer_push_constants), {1, 1, 1}, {(uint32_t)k_type, fa_block_bytes(k_type), device->subgroup_size}, 1, true, li_subgroup);
}
}
{
const uint32_t gdn_sizes[] = {16, 32, 64, 128};
const char * gdn_names[][2] = {
@@ -11697,6 +11749,12 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const
return ctx->device->pipeline_gated_linear_attn_f32;
}
return nullptr;
case GGML_OP_LIGHTNING_INDEXER:
// only the k type selects a pipeline, the other types are fixed by ggml_lightning_indexer()
if (ggml_vk_lightning_indexer_k_type_supported(src1->type)) {
return ctx->device->pipeline_lightning_indexer_f32[src1->type];
}
return nullptr;
case GGML_OP_GATED_DELTA_NET:
if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
const uint32_t S_v = dst->src[2]->ne[0];
@@ -12772,6 +12830,55 @@ static void ggml_vk_gated_linear_attn(ggml_backend_vk_context * ctx, vk_context&
pc, { (uint32_t)(n_seqs * n_heads), 1, 1 });
}
static void ggml_vk_lightning_indexer(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst) {
const ggml_tensor * q = dst->src[0];
const ggml_tensor * k = dst->src[1];
const ggml_tensor * w = dst->src[2];
const ggml_tensor * m = dst->src[3];
vk_pipeline pipeline = ggml_vk_op_get_pipeline(ctx, q, k, w, dst, dst->op);
GGML_ASSERT(pipeline != nullptr);
ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1);
const uint32_t n_kv = k->ne[2];
const uint32_t n_heads = q->ne[1];
const uint32_t n_tokens = q->ne[2];
const uint32_t n_streams = q->ne[3];
const uint32_t n_masks = m->ne[3];
const uint32_t n_outputs = (uint32_t)(dst->ne[0] * dst->ne[1] * dst->ne[3]);
const uint32_t dispatch_x = std::min(n_outputs, ctx->device->properties.limits.maxComputeWorkGroupCount[0]);
const uint32_t dispatch_y = CEIL_DIV(n_outputs, dispatch_x);
// q, w and dst are f32 and m is f16, so their strides are passed in elements;
// k may be quantized, so its strides stay in bytes
const uint32_t q_nb1 = q->nb[1] / sizeof(float);
const uint32_t q_nb2 = q->nb[2] / sizeof(float);
const uint32_t q_nb3 = q->nb[3] / sizeof(float);
const uint32_t k_nb2 = k->nb[2];
const uint32_t k_nb3 = k->nb[3];
const uint32_t w_nb1 = w->nb[1] / sizeof(float);
const uint32_t w_nb3 = w->nb[3] / sizeof(float);
const uint32_t m_nb1 = m->nb[1] / sizeof(ggml_fp16_t);
const uint32_t m_nb3 = m->nb[3] / sizeof(ggml_fp16_t);
const uint32_t d_nb1 = dst->nb[1] / sizeof(float);
const uint32_t d_nb3 = dst->nb[3] / sizeof(float);
const vk_op_lightning_indexer_push_constants pc = {
n_kv, n_heads, n_tokens, n_streams, n_masks, dispatch_x,
q_nb1, q_nb2, q_nb3,
k_nb2, k_nb3,
w_nb1, w_nb3,
m_nb1, m_nb3,
d_nb1, d_nb3,
};
ggml_vk_dispatch_pipeline(ctx, subctx, pipeline,
{ggml_vk_tensor_subbuffer(ctx, q), ggml_vk_tensor_subbuffer(ctx, k), ggml_vk_tensor_subbuffer(ctx, w), ggml_vk_tensor_subbuffer(ctx, m), ggml_vk_tensor_subbuffer(ctx, dst)},
pc, {dispatch_x, dispatch_y, 1});
}
static void ggml_vk_gated_delta_net(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst) {
const ggml_tensor * src_q = dst->src[0];
const ggml_tensor * src_v = dst->src[2];
@@ -15898,6 +16005,11 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr
break;
case GGML_OP_LIGHTNING_INDEXER:
ggml_vk_lightning_indexer(ctx, compute_ctx, node);
break;
case GGML_OP_GATED_DELTA_NET:
ggml_vk_gated_delta_net(ctx, compute_ctx, node);
@@ -18676,6 +18788,40 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
case GGML_OP_GATED_LINEAR_ATTN:
// the shader block size is hardcoded to head_size 64
return op->src[0]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32 && op->src[0]->ne[0] == 64;
case GGML_OP_LIGHTNING_INDEXER:
{
const ggml_tensor * q = op->src[0];
const ggml_tensor * k = op->src[1];
const ggml_tensor * w = op->src[2];
const ggml_tensor * m = op->src[3];
// the q/w/m types and the shape relationships between q, k, w, m and dst
// are already asserted in ggml_lightning_indexer()
if (!ggml_vk_lightning_indexer_k_type_supported(k->type) || !device->fp16) {
return false;
}
// the shader block size is hardcoded to head size 128
if (q->ne[0] != 128) {
return false;
}
// the shader indexes the buffers by element stride, and is dispatched
// without allow_misalign
for (const ggml_tensor * t : {q, k, w, m, op}) {
if (t->nb[0] != ggml_type_size(t->type) ||
(vk_tensor_offset(t) + t->view_offs) % device->properties.limits.minStorageBufferOffsetAlignment != 0) {
return false;
}
// the strides get scaled down from bytes, so the division must be exact
for (int i = 1; i < GGML_MAX_DIMS; ++i) {
if (t->nb[i] % ggml_type_size(t->type) != 0) {
return false;
}
}
}
return true;
}
case GGML_OP_GATED_DELTA_NET:
{
const uint32_t S_v = op->src[2]->ne[0];
@@ -19685,6 +19831,8 @@ static void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph *
const float * op_params = (const float *)tensor->op_params;
tensor_clone = ggml_gated_linear_attn(ggml_ctx, src_clone[0], src_clone[1],
src_clone[2], src_clone[3], src_clone[4], op_params[0]);
} else if (tensor->op == GGML_OP_LIGHTNING_INDEXER) {
tensor_clone = ggml_lightning_indexer(ggml_ctx, src_clone[0], src_clone[1], src_clone[2], src_clone[3]);
} else if (tensor->op == GGML_OP_GATED_DELTA_NET) {
tensor_clone = ggml_gated_delta_net(ggml_ctx, src_clone[0], src_clone[1],
src_clone[2], src_clone[3], src_clone[4], src_clone[5],
@@ -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)
@@ -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.
@@ -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]);
}
}
@@ -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"}}));