Add fused DeepSeek4 optimization ops#186
Conversation
Assisted-by: GPT-5.6 Sol
Assisted-by: GPT-5.6 Sol
Assisted-by: GPT-5.6 Sol
Assisted-by: GPT-5.6 Sol
Assisted-by: GPT-5.6 Sol
Assisted-by: GPT-5.6 Sol
Assisted-by: GPT-5.6 Sol
Assisted-by: GPT-5.6 Sol
Assisted-by: GPT-5.6 Sol
Assisted-by: GPT-5.6 Sol
Assisted-by: GPT-5.6 Sol
Assisted-by: GPT-5.6 Sol
Assisted-by: GPT-5.6 Sol
Assisted-by: GPT-5.6 Sol
Assisted-by: GPT-5.6 Sol
| indexer_score = ggml_mul(ctx0, indexer_score, indexer_weights); | ||
| indexer_score = ggml_sum_rows(ctx0, indexer_score); | ||
| indexer_score = ggml_cont(ctx0, ggml_permute(ctx0, indexer_score, 2, 1, 0, 3)); | ||
| ggml_tensor * indexer_score = ggml_dsv4_lid_score(ctx0, indexer_kq, indexer_weights); |
There was a problem hiding this comment.
The matmul result indexer_kq is still fully materialized before ggml_dsv4_lid_score consumes it. Could we investigate a DeepSeek4-specific matmul epilogue that applies ReLU, multiplies by indexer_weights, and reduces over heads while accumulators are still resident? This should eliminate the largest LID intermediate, plus its layout/contiguous conversion and an extra kernel launch. The current op can remain as the fallback.
| const int64_t d = (ir / n_embd) % hc; | ||
| const int64_t it = ir / (n_embd * hc); | ||
|
|
||
| volatile float sum = (*(const float *) ((const char *) x->data + i*x->nb[0] + it*x->nb[1])) * |
There was a problem hiding this comment.
Could we benchmark this without volatile (or isolate the exact bit-exact requirement)? volatile on the accumulator and every inner-loop product can inhibit register allocation/vectorization and force unnecessary loads/stores on CPU. If strict ordering is required for backend bit-exact tests, a non-volatile fast path with an explicit deterministic reduction may preserve reproducibility at much lower cost.
| case GGML_OP_DSV4_LID_SCORE: | ||
| { | ||
| ggml_compute_forward_dsv4_lid_score(params, tensor); | ||
| } break; | ||
| case GGML_OP_DSV4_HC_POST: | ||
| { | ||
| ggml_compute_forward_dsv4_hc_post(params, tensor); | ||
| } break; | ||
| case GGML_OP_DSV4_HC_SINKHORN: |
There was a problem hiding this comment.
Three new model-specific ops in the core ggml op enum. AGENTS.md explicitly says to "prefer reusing existing infrastructure over introducing new components" and to pause on changes that "add whole new subsystems." Adding DSV4_* (DeepSeek4-specific) entries to the global GGML_OP_* enum, RPC protocol, backend-meta, and all four backends is a large, invasive, model-specific addition. Upstream ggml generally resists model-specific fused ops in the core op set.
Could be worthwhile to consider if a more generic shape can be found for the ops. For example, DSV4_HC_SINKHORN is the most clearly generalizable.
Sinkhorn-Knopp normalization (alternating row/column normalization toward a doubly-stochastic matrix) is a standard, model-agnostic algorithm. It shows up in optimal-transport layers and in some MoE routing schemes. Nothing about the algorithm is DeepSeek4-specific - only the implementation is, in two ways:
- It's hardcoded to 4x4 (16 registers, float4 reductions, fully unrolled loops).
- The op takes eps and n_iter as fixed params, which is already generic.
A clean generalization would be ggml_sinkhorn(x, eps, n_iter) operating on the leading NxM of each [N, M, batch, 1] tensor, with the kernel special-casing small sizes (4x4, 8x8) for the register-resident fast path and falling back to shared memory for larger matrices. That would be a defensible standalone ggml op that other models could reuse.
DSV4_HC_POST is a batched small GEMM + rank-1 update, also somewhat generic and could at least be renamed.
Assisted-by: GPT-5.6 Sol
Assisted-by: GPT-5.6 Sol
Assisted-by: GPT-5.6 Sol
Assisted-by: GPT-5.6 Sol
Assisted-by: GPT-5.6 Sol
|
Upstream already implemented some of this PR. Looking at the next rebase target
So I think this PR should instead target the temp-10069, and fill the gaps with the Metal kernels and Vulkan shaders. Otherwise, we will need to carry on competitive implementations. |
Summary
test-backend-opscoverage for the new ops.Testing
git diff --check tether/temp-9840...HEADcmake -S . -B build-review -DLLAMA_BUILD_TESTS=ON -DLLAMA_BUILD_TOOLS=OFF -DLLAMA_BUILD_EXAMPLES=OFF -DLLAMA_BUILD_SERVER=OFF -DLLAMA_BUILD_APP=OFF -DGGML_METAL=OFF -DGGML_VULKAN=OFF -DGGML_CUDA=OFFcmake --build build-review --target test-backend-ops -j 8./build-review/bin/test-backend-ops -b CPU -o DSV4_LID_SCORE_BIT_EXACT,DSV4_HC_POST_BIT_EXACT,DSV4_HC_SINKHORN_BIT_EXACT(15/15 passed)cmake -S . -B build-review-metal -DLLAMA_BUILD_TESTS=ON -DLLAMA_BUILD_TOOLS=OFF -DLLAMA_BUILD_EXAMPLES=OFF -DLLAMA_BUILD_SERVER=OFF -DLLAMA_BUILD_APP=OFF -DGGML_METAL=ON -DGGML_VULKAN=OFF -DGGML_CUDA=OFFcmake --build build-review-metal --target test-backend-ops -j 8./build-review-metal/bin/test-backend-ops -b MTL0 -o DSV4_LID_SCORE_BIT_EXACT,DSV4_HC_POST_BIT_EXACT,DSV4_HC_SINKHORN_BIT_EXACT(15/15 passed on Apple M4)