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/ggml-vulkan.cpp
+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],