Added LIGHTNING_INDEXER support for Deepseek V4 ops on Vulkan Backend - #27008
Added LIGHTNING_INDEXER support for Deepseek V4 ops on Vulkan Backend#27008shenron0101 wants to merge 1 commit into
Conversation
github-actions
Bot
added
testing
|
Hi @shenron0101, thanks for your contribution! Per our contribution guidelines, the automated PR checker found the following issue(s) that need your attention:
Please note that maintainers reserve the right to make final decisions on PRs. If you believe there is a mistake, please comment below. |
| } | ||
|
|
||
| if (tid == 0) { | ||
| score += max(partials[0], 0.0) * weights[h + t * w_nb1 + s * w_nb3]; |
There was a problem hiding this comment.
I think this will race against line 103 on the next loop iteration. I also think line 108 will race against itself. You can run this against the vulkan validation layers sharedmemorydatarace checker to verify the fixes.
| const uint output_idx = gl_WorkGroupID.y * dispatch_x + gl_WorkGroupID.x; | ||
| const uint n_outputs = n_kv * n_tokens * n_streams; | ||
|
|
||
| if (K_TYPE == FA_TYPE_IQ4_NL) { |
There was a problem hiding this comment.
Would be nice to reuse fa_type_needs_shmem(K_TYPE). Might need to move that and FA_TYPE_ definitions into a new header.
| q->type != GGML_TYPE_F32 || | ||
| !ggml_vk_lightning_indexer_k_type_supported(k->type) || | ||
| w->type != GGML_TYPE_F32 || m->type != GGML_TYPE_F16 || op->type != GGML_TYPE_F32 || | ||
| !device->fp16 || device->properties.limits.maxComputeWorkGroupInvocations < 128 || |
There was a problem hiding this comment.
some of these checks aren't necessary due to guaranteed spec minimums.
| return false; | ||
| } | ||
|
|
||
| if (q->ne[0] != 128 || k->ne[0] != 128 || k->ne[1] != 1 || |
There was a problem hiding this comment.
This whole if test is unreadable and should be refactored somehow. I also think some of it is redundant with asserts in ggml_lightning_indexer and not needed here.
| if (tensor->ne[0] <= 0 || tensor->ne[1] <= 0 || tensor->ne[2] <= 0 || tensor->ne[3] <= 0 || | ||
| tensor->nb[0] != ggml_type_size(tensor->type) || | ||
| !storage_buffer_offset_aligned(tensor) || | ||
| ggml_nbytes(tensor) > device->properties.limits.maxStorageBufferRange) { |
There was a problem hiding this comment.
maxStorageBufferRange is already checked. And the <=0 checks aren't needed.
There was a problem hiding this comment.
Wow, there's a lot more after this. Please reconsider these checks to be more consistent with others.
| #include "ggml-alloc.h" | ||
| #include "ggml-backend.h" | ||
| #include "ggml-cpp.h" | ||
| #include "../ggml/src/ggml-impl.h" |
fb1107d to
10bb8f1
Compare
Implements GGML_OP_LIGHTNING_INDEXER for the Vulkan backend, with K accepted as f32, f16, bf16, q8_0, q5_1, q5_0, q4_1, q4_0 and iq4_nl. One workgroup of 128 invocations computes one output element: it stages the K row in shared memory, then loops over heads accumulating the ReLU'd q.k score weighted by the per-head weight, and adds the mask. The FA_TYPE_* spec-constant values and their helpers move out of flash_attn_base.glsl into a new fa_types.glsl so the indexer shader can share them instead of redeclaring them. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
10bb8f1 to
7de70a6
Compare
|
Intentional or accident? |
|
Sorry, I was using a fork of a fork which made development tricky. I ported over the changes to a new branch and will draft a new PR shortly. |
3 participants
Overview
LIGHTNING_INDEXER: Added support for F32, F16, BF16, and several quantized formats.
The Lightning Indexer has only been implemented in the CPU backend, CUDA backend but not in the Vulkan backend. This is an implementation of the Indexer for DeepSeek V4 on the Vulkan backend.
Assisted-by: OpenCode
Currently the Hyper-connection Operations for the Vulkan backend are being implemented by @kh0pper. The one missing operation in the PR was the lightning indexer.
#26585 — "vulkan: tiled transpose for 0<->2 permuted CONT"
The original PR which was closed:
#26548 — "vulkan: add DeepSeek-V4 hyper-connection fused ops (DSV4_HC_COMB/PRE/POST)".
Additional information
Vulkan Lightning Indexer performance (RTX 3090)
Correctness: all 171 test-backend-ops cases pass on Vulkan0 — every K-type (F32/F16/BF16/Q8_0/Q5_1/Q5_0/Q4_1/Q4_0/IQ4_NL) and every masked/padded variant.
Throughput, full sweep (324 shapes = 3 KV lengths × 3 batch sizes × 2 head counts × 2 stream counts × 9 K-types):
┌───────────┬────────────┬───────┬───────┐
│ KV length │ avg GFLOPS │ min │ max │
├───────────┼────────────┼───────┼───────┤
│ 256 │ 163.0 │ 111.2 │ 177.9 │
├───────────┼────────────┼───────┼───────┤
│ 4,096 │ 174.4 │ 169.4 │ 177.0 │
├───────────┼────────────┼───────┼───────┤
│ 65,536 │ 175.2 │ 171.4 │ 178.2 │
└───────────┴────────────┴───────┴───────┘
Steady-state throughput lands around ~171–178 GFLOPS regardless of K quantization format and the quantized formats aren't leaving meaningful performance on the table versus F32/F16/BF16.
Requirements