Skip to content

perf(ds4): enable causal Flash prefill #512

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

Open
weicj wants to merge 10 commits into
base: main
Choose a base branch
from
Open

perf(ds4): enable causal Flash prefill #512

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
b9c0069
feat(ggml): vendor ROCmFPX quant support for DeepSeek V4 Flash
davide221 Jul 8, 2026
a5f1146
fix(ggml): make ROCmFP4 codebook self-contained on CUDA
davide221 Jul 8, 2026
42f1e10
fix(deepseek4): ROCmFPX byte-safe decode + fused-decode path
davide221 Jul 8, 2026
6c7ee19
feat(deepseek4): add opt-in expert top-k knobs
Jul 9, 2026
b6fcc26
feat(dflash): bridge DeepSeek4 DSpark aux heads
davide221 Jul 7, 2026
b989d3f
feat(deepseek4): add DSpark speculative runtime
davide221 Jul 10, 2026
867c9b8
perf(deepseek4): use DSpark confidence for adaptive verify width
davide221 Jul 10, 2026
95c9812
perf(deepseek4): enable causal Flash prefill
Jul 11, 2026
c46e945
fix(deepseek4): harden causal Flash prefill
Jul 12, 2026
0771934
fix(deepseek4): reset weights after failed reload
Jul 12, 2026
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
18 changes: 18 additions & 0 deletions server/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -258,6 +258,8 @@ add_library(dflash_common STATIC
src/deepseek4/deepseek4_daemon.cpp
src/deepseek4/deepseek4_layer_split_adapter.cpp
src/deepseek4/deepseek4_target_shard_ipc_daemon.cpp
src/deepseek4/deepseek4_dspark.cpp
src/deepseek4/deepseek4_dspark_spec.cpp
src/flashprefill_q8.cpp
src/kv_cache.cpp
src/kv_quant.cpp
Expand Down Expand Up @@ -801,6 +803,22 @@ if(DFLASH27B_TESTS)
endif()
add_test(NAME deepseek4_unit COMMAND test_deepseek4_unit)
endif()
if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_ds4_dspark_load.cpp")
add_executable(test_ds4_dspark_load test/test_ds4_dspark_load.cpp)
target_include_directories(test_ds4_dspark_load PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS} ${CMAKE_CURRENT_SOURCE_DIR}/include)
if(DFLASH27B_GPU_BACKEND STREQUAL "hip")
target_compile_definitions(test_ds4_dspark_load PRIVATE DFLASH27B_BACKEND_HIP=1 GGML_USE_HIP)
else()
target_compile_definitions(test_ds4_dspark_load PRIVATE DFLASH27B_BACKEND_CUDA=1 GGML_USE_CUDA)
endif()
target_link_libraries(test_ds4_dspark_load PRIVATE dflash_common ggml ggml-cpu ${DFLASH27B_GGML_BACKEND_TARGET} ggml-base)
if(DFLASH27B_GPU_BACKEND STREQUAL "cuda")
find_package(CUDAToolkit REQUIRED)
target_link_libraries(test_ds4_dspark_load PRIVATE CUDA::cudart)
else()
target_link_libraries(test_ds4_dspark_load PRIVATE hip::host)
endif()
endif()
if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/smoke_load_draft.cpp")
add_executable(smoke_load_draft test/smoke_load_draft.cpp)
target_include_directories(smoke_load_draft PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS})
Expand Down
60 changes: 59 additions & 1 deletion server/deps/llama.cpp/ggml/include/ggml.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -430,7 +430,13 @@ extern "C" {
GGML_TYPE_NVFP4 = 40, // NVFP4 (4 blocks, E4M3 scale)
GGML_TYPE_Q1_0 = 41,
GGML_TYPE_TQ3_0 = 42, // TurboQuant 3.5 bpv (3-bit Lloyd-Max + FWHT rotation)
GGML_TYPE_COUNT = 43,
GGML_TYPE_Q4_0_ROCMFP4 = 100,
GGML_TYPE_Q4_0_ROCMFP4_FAST = 101,
GGML_TYPE_Q6_0_ROCMFPX = 102,
GGML_TYPE_Q8_0_ROCMFPX = 103,
GGML_TYPE_Q3_0_ROCMFPX = 104,
GGML_TYPE_Q2_0_ROCMFP2 = 107,
GGML_TYPE_COUNT = 108,
};

// precision
Expand Down Expand Up @@ -468,6 +474,18 @@ extern "C" {
GGML_FTYPE_MOSTLY_MXFP4 = 25, // except 1d tensors
GGML_FTYPE_MOSTLY_NVFP4 = 26, // except 1d tensors
GGML_FTYPE_MOSTLY_Q1_0 = 27, // except 1d tensors
GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4 = 100,
GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4_LEAN = 101,
GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4_COHERENT = 102,
GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4_FAST = 103,
GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4_FAST_COHERENT = 104,
GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4_STRIX = 105,
GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4_STRIX_LEAN = 106,
GGML_FTYPE_MOSTLY_Q6_0_ROCMFPX = 110,
GGML_FTYPE_MOSTLY_Q8_0_ROCMFPX = 111,
GGML_FTYPE_MOSTLY_Q3_0_ROCMFPX = 112,
GGML_FTYPE_MOSTLY_Q2_0_ROCMFP2 = 118,
GGML_FTYPE_MOSTLY_Q2_0_ROCMFP2_STRIX = 119,
};

// available tensor operations:
Expand Down Expand Up @@ -583,6 +601,8 @@ extern "C" {

GGML_OP_MOE_FUSED, // Fused MoE FFN: gate+up+swiglu+down+weighted_sum+shared_expert

GGML_OP_DS4_HC, // Fused DeepSeek4 hyper-connection pre/post/out mixing

GGML_OP_COUNT,
};

Expand Down Expand Up @@ -618,6 +638,7 @@ extern "C" {
GGML_GLU_OP_GEGLU,
GGML_GLU_OP_SWIGLU,
GGML_GLU_OP_SWIGLU_OAI,
GGML_GLU_OP_SWIGLU_DS4,
GGML_GLU_OP_GEGLU_ERF,
GGML_GLU_OP_GEGLU_QUICK,

Expand Down Expand Up @@ -1339,6 +1360,12 @@ extern "C" {
struct ggml_tensor * a,
struct ggml_tensor * b);

GGML_API struct ggml_tensor * ggml_swiglu_ds4_split(
struct ggml_context * ctx,
struct ggml_tensor * gate,
struct ggml_tensor * up,
float clamp);

GGML_API struct ggml_tensor * ggml_geglu_erf_split(
struct ggml_context * ctx,
struct ggml_tensor * a,
Expand Down Expand Up @@ -2399,6 +2426,37 @@ extern "C" {
struct ggml_tensor * experts,
struct ggml_tensor * expert_weights);

// Fused DeepSeek4 hyper-connection helpers (decode, n_tokens == 1).
// ggml_ds4_hc_pre: mix[2*n_hc+n_hc^2] + base + hc_state[n_embd*n_hc] ->
// dst[n_embd + 2*n_hc + n_hc^2] = { working, split(pre,post,comb) }
GGML_API struct ggml_tensor * ggml_ds4_hc_pre(
struct ggml_context * ctx,
struct ggml_tensor * mix,
struct ggml_tensor * base,
struct ggml_tensor * hc_state,
int n_hc,
int sinkhorn_iters,
float pre_scale,
float post_scale,
float comb_scale);

// ggml_ds4_hc_post: residual hc_state + block_out + split -> new hc_state
GGML_API struct ggml_tensor * ggml_ds4_hc_post(
struct ggml_context * ctx,
struct ggml_tensor * residual_hc,
struct ggml_tensor * block_out,
struct ggml_tensor * split,
int n_hc);

// ggml_ds4_hc_out: output-stage merge of hc streams into one embedding
GGML_API struct ggml_tensor * ggml_ds4_hc_out(
struct ggml_context * ctx,
struct ggml_tensor * mix,
struct ggml_tensor * base,
struct ggml_tensor * hc_state,
int n_hc,
float pre_scale);

// TODO: needs to be adapted to ggml_flash_attn_ext
GGML_API struct ggml_tensor * ggml_flash_attn_back(
struct ggml_context * ctx,
Expand Down
Loading
Loading