Skip to content

Added LIGHTNING_INDEXER support for Deepseek V4 ops on Vulkan Backend #27008

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

Closed
shenron0101 wants to merge 1 commit into from
Closed

Added LIGHTNING_INDEXER support for Deepseek V4 ops on Vulkan Backend #27008

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
136 changes: 136 additions & 0 deletions ggml/src/ggml-vulkan/ggml-vulkan.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -1063,6 +1078,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;
Expand Down Expand Up @@ -1842,6 +1858,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;
Expand Down Expand Up @@ -5803,6 +5839,12 @@ 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);

for (ggml_type k_type : lightning_indexer_k_types) {
const std::string name = "lightning_indexer_" + std::string(ggml_type_name(k_type)) + "_k_f32";
const uint32_t block_bytes = k_type == GGML_TYPE_F32 ? 4 * sizeof(float) : (uint32_t)ggml_type_size(k_type);
ggml_vk_create_pipeline(device, device->pipeline_lightning_indexer_f32[k_type], name.c_str(), lightning_indexer_f32_len, lightning_indexer_f32_data, "main", 5, sizeof(vk_op_lightning_indexer_push_constants), {1, 1, 1}, {(uint32_t)k_type, block_bytes}, 1);
}

{
const uint32_t gdn_sizes[] = {16, 32, 64, 128};
const char * gdn_names[][2] = {
Expand Down Expand Up @@ -11554,6 +11596,13 @@ 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:
if (src0->type == GGML_TYPE_F32 && src2->type == GGML_TYPE_F32 && dst->src[3]->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F32) {
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];
Expand Down Expand Up @@ -12619,6 +12668,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];
Expand Down Expand Up @@ -15623,6 +15721,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);

Expand Down Expand Up @@ -18385,6 +18488,37 @@ 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];
if (!q || !k || !w || !m) {
return false;
}

// 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 || k->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;
}
}
return true;
}
case GGML_OP_GATED_DELTA_NET:
{
const uint32_t S_v = op->src[2]->ne[0];
Expand Down Expand Up @@ -19378,6 +19512,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],
Expand Down
50 changes: 50 additions & 0 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
@@ -0,0 +1,50 @@
// 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;
}
}
51 changes: 1 addition & 50 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 @@ -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
Expand All @@ -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.
Expand Down
21 changes: 20 additions & 1 deletion 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 @@ -15,23 +15,39 @@
// F32 is fed as a vec4 "block" (4 floats), matching what dequant_funcs_cm2.glsl
// does for F32 in the cm2 shader. FaBlockBytesK/V == 16 for F32.
layout (binding = 1) readonly buffer K_PACKED_F32 { vec4 data[]; } k_packed_f32;
#ifndef FA_K_ONLY
layout (binding = 2) readonly buffer V_PACKED_F32 { vec4 data[]; } v_packed_f32;
#endif

layout (binding = 1) readonly buffer K_PACKED_Q4_0 { block_q4_0_packed16 data[]; } k_packed_q4_0;
#ifndef FA_K_ONLY
layout (binding = 2) readonly buffer V_PACKED_Q4_0 { block_q4_0_packed16 data[]; } v_packed_q4_0;
#endif
layout (binding = 1) readonly buffer K_PACKED_Q4_1 { block_q4_1_packed16 data[]; } k_packed_q4_1;
#ifndef FA_K_ONLY
layout (binding = 2) readonly buffer V_PACKED_Q4_1 { block_q4_1_packed16 data[]; } v_packed_q4_1;
#endif
layout (binding = 1) readonly buffer K_PACKED_Q5_0 { block_q5_0_packed16 data[]; } k_packed_q5_0;
#ifndef FA_K_ONLY
layout (binding = 2) readonly buffer V_PACKED_Q5_0 { block_q5_0_packed16 data[]; } v_packed_q5_0;
#endif
layout (binding = 1) readonly buffer K_PACKED_Q5_1 { block_q5_1_packed16 data[]; } k_packed_q5_1;
#ifndef FA_K_ONLY
layout (binding = 2) readonly buffer V_PACKED_Q5_1 { block_q5_1_packed16 data[]; } v_packed_q5_1;
#endif
layout (binding = 1) readonly buffer K_PACKED_Q8_0 { block_q8_0_packed16 data[]; } k_packed_q8_0;
#ifndef FA_K_ONLY
layout (binding = 2) readonly buffer V_PACKED_Q8_0 { block_q8_0_packed16 data[]; } v_packed_q8_0;
#endif
layout (binding = 1) readonly buffer K_PACKED_IQ4_NL { block_iq4_nl_packed16 data[]; } k_packed_iq4_nl;
#ifndef FA_K_ONLY
layout (binding = 2) readonly buffer V_PACKED_IQ4_NL { block_iq4_nl_packed16 data[]; } v_packed_iq4_nl;
#endif

layout (binding = 1) readonly buffer K_PACKED_BF16 { u16vec4 data[]; } k_packed_bf16;
#ifndef FA_K_ONLY
layout (binding = 2) readonly buffer V_PACKED_BF16 { u16vec4 data[]; } v_packed_bf16;
#endif

// Q4_1 and Q5_1 packed32 views: aliased to the same memory as the packed16
// views, used by the MMQ K-side hot path for fast 4-uint loads.
Expand Down Expand Up @@ -130,7 +146,9 @@ FLOAT_TYPEV4 dequantize4(uint ib, uint iqs, uint a_offset, uint binding_idx) {
case FA_TYPE_IQ4_NL: FA_DEQUANT4_IQ4_NL(k_packed_iq4_nl)
case FA_TYPE_BF16: FA_DEQUANT4_BF16(k_packed_bf16)
}
} else {
}
#ifndef FA_K_ONLY
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)
Expand All @@ -142,5 +160,6 @@ FLOAT_TYPEV4 dequantize4(uint ib, uint iqs, uint a_offset, uint binding_idx) {
case FA_TYPE_BF16: FA_DEQUANT4_BF16(v_packed_bf16)
}
}
#endif
return FLOAT_TYPEV4(0);
}
Loading